From 57b12a7aca99097d42278f174bf3919ed6c9cf08 Mon Sep 17 00:00:00 2001 From: Semih702 Date: Tue, 10 Feb 2026 01:58:42 +0300 Subject: [PATCH] ci: add Go unit tests and enforce go quality checks (race-enabled) --- .github/workflows/go-ci.yaml | 4 +- README.md | 4 +- collector/Dockerfile | 41 +- collector/cmd/collector/main.go | 30 ++ collector/go.mod | 8 + collector/go.sum | 10 + collector/internal/collector/server.go | 102 ++++ collector/internal/collector/server_test.go | 61 +++ collector/internal/collector/types.go | 18 + collector/internal/collector/util.go | 11 + collector/main.go | 120 ----- proxy/Dockerfile | 42 +- proxy/cmd/proxy/main.go | 32 ++ proxy/go.mod | 8 + proxy/go.sum | 10 + proxy/internal/proxy/capture.go | 33 ++ proxy/internal/proxy/capture_test.go | 22 + proxy/internal/proxy/config.go | 28 + proxy/internal/proxy/handler.go | 288 ++++++++++ proxy/internal/proxy/sse.go | 62 +++ proxy/internal/proxy/sse_test.go | 39 ++ proxy/internal/proxy/types.go | 53 ++ proxy/internal/proxy/util.go | 91 ++++ proxy/internal/proxy/util_test.go | 23 + proxy/main.go | 560 -------------------- 25 files changed, 975 insertions(+), 725 deletions(-) create mode 100644 collector/cmd/collector/main.go create mode 100644 collector/go.sum create mode 100644 collector/internal/collector/server.go create mode 100644 collector/internal/collector/server_test.go create mode 100644 collector/internal/collector/types.go create mode 100644 collector/internal/collector/util.go delete mode 100644 collector/main.go create mode 100644 proxy/cmd/proxy/main.go create mode 100644 proxy/go.sum create mode 100644 proxy/internal/proxy/capture.go create mode 100644 proxy/internal/proxy/capture_test.go create mode 100644 proxy/internal/proxy/config.go create mode 100644 proxy/internal/proxy/handler.go create mode 100644 proxy/internal/proxy/sse.go create mode 100644 proxy/internal/proxy/sse_test.go create mode 100644 proxy/internal/proxy/types.go create mode 100644 proxy/internal/proxy/util.go create mode 100644 proxy/internal/proxy/util_test.go delete mode 100644 proxy/main.go diff --git a/.github/workflows/go-ci.yaml b/.github/workflows/go-ci.yaml index 11e168a..b665cb3 100644 --- a/.github/workflows/go-ci.yaml +++ b/.github/workflows/go-ci.yaml @@ -65,8 +65,8 @@ jobs: set -euo pipefail go vet ./... - - name: go test + - name: go test (race) working-directory: ${{ matrix.module }} run: | set -euo pipefail - go test ./... + go test -race ./... diff --git a/README.md b/README.md index 6ec5bc9..578fffc 100644 --- a/README.md +++ b/README.md @@ -372,8 +372,8 @@ This pipeline exists to ensure the Helm chart remains: - Add CI checks for Go code quality: - `go test ./...` (with `-race` where feasible) - `golangci-lint` (or at least `go vet`) - - formatting checks (`gofmt`) and module tidiness (`go mod tidy` / `go mod verify`) -- Keep these checks fast to run on every PR, and required before release + - formatting checks (`gofmt`) and module tidiness (`go mod tidy` / `go mod verify`) ✅ +- Keep these checks fast to run on every PR, and required before release ✅ ### Streaming (SSE) – final validation (edge cases) - Client disconnect propagation (client → proxy → upstream) diff --git a/collector/Dockerfile b/collector/Dockerfile index 5827465..843ca4b 100644 --- a/collector/Dockerfile +++ b/collector/Dockerfile @@ -1,20 +1,21 @@ -# -------- build stage -------- -FROM golang:1.22-alpine AS builder - -WORKDIR /src -COPY go.mod ./ -RUN go mod download - -COPY . . -RUN CGO_ENABLED=0 GOOS=linux GOARCH=amd64 go build -trimpath -ldflags="-s -w" -o /out/llm-collector . - -# -------- runtime stage -------- -FROM gcr.io/distroless/static:nonroot - -WORKDIR / -COPY --from=builder /out/llm-collector /llm-collector - -EXPOSE 8081 -USER nonroot:nonroot - -ENTRYPOINT ["/llm-collector"] +# -------- build stage -------- +FROM golang:1.22-alpine AS builder + +WORKDIR /src +COPY go.mod go.sum ./ +RUN go mod download + +COPY . . +RUN CGO_ENABLED=0 GOOS=linux GOARCH=amd64 \ + go build -trimpath -ldflags="-s -w" -o /out/llm-collector ./cmd/collector + +# -------- runtime stage -------- +FROM gcr.io/distroless/static:nonroot + +WORKDIR / +COPY --from=builder /out/llm-collector /llm-collector + +EXPOSE 8081 +USER nonroot:nonroot + +ENTRYPOINT ["/llm-collector"] diff --git a/collector/cmd/collector/main.go b/collector/cmd/collector/main.go new file mode 100644 index 0000000..2bf3d78 --- /dev/null +++ b/collector/cmd/collector/main.go @@ -0,0 +1,30 @@ +package main + +import ( + "log" + "net/http" + "time" + + collector "llm-collector/internal/collector" +) + +func main() { + addr := ":" + collector.Getenv("PORT", "8081") + outPath := collector.Getenv("EVENT_LOG_PATH", "") + + s, err := collector.NewServer(outPath) + if err != nil { + log.Fatalf("open log file: %v", err) + } + defer s.Close() + + log.Printf("collector listening on %s", addr) + + srv := &http.Server{ + Addr: addr, + Handler: s.Mux(), + ReadHeaderTimeout: 5 * time.Second, + } + + log.Fatal(srv.ListenAndServe()) +} diff --git a/collector/go.mod b/collector/go.mod index 052858c..1a6c4f3 100644 --- a/collector/go.mod +++ b/collector/go.mod @@ -1,3 +1,11 @@ module llm-collector go 1.22 + +require github.com/stretchr/testify v1.9.0 + +require ( + github.com/davecgh/go-spew v1.1.1 // indirect + github.com/pmezard/go-difflib v1.0.0 // indirect + gopkg.in/yaml.v3 v3.0.1 // indirect +) diff --git a/collector/go.sum b/collector/go.sum new file mode 100644 index 0000000..60ce688 --- /dev/null +++ b/collector/go.sum @@ -0,0 +1,10 @@ +github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= +github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= +github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= +github.com/stretchr/testify v1.9.0 h1:HtqpIVDClZ4nwg75+f6Lvsy/wHu+3BoSGCbBAcpTsTg= +github.com/stretchr/testify v1.9.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY= +gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM= +gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= +gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= +gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= diff --git a/collector/internal/collector/server.go b/collector/internal/collector/server.go new file mode 100644 index 0000000..7b01303 --- /dev/null +++ b/collector/internal/collector/server.go @@ -0,0 +1,102 @@ +package collector + +import ( + "bufio" + "encoding/json" + "log" + "net/http" + "os" + "sync" + "time" +) + +type Server struct { + mu sync.Mutex + file *os.File + w *bufio.Writer +} + +// NewServer creates a server. If outPath is empty, events are printed to stdout. +func NewServer(outPath string) (*Server, error) { + s := &Server{} + if outPath == "" { + log.Printf("collector: EVENT_LOG_PATH not set; events will be printed to stdout") + return s, nil + } + + f, err := os.OpenFile(outPath, os.O_CREATE|os.O_APPEND|os.O_WRONLY, 0644) + if err != nil { + return nil, err + } + + s.file = f + s.w = bufio.NewWriterSize(f, 1<<20) + log.Printf("collector: writing events to %s", outPath) + return s, nil +} + +// Close flushes and closes the underlying file if configured. +func (s *Server) Close() { + s.mu.Lock() + defer s.mu.Unlock() + + if s.w != nil { + _ = s.w.Flush() + } + if s.file != nil { + _ = s.file.Close() + } +} + +func (s *Server) Mux() *http.ServeMux { + mux := http.NewServeMux() + mux.HandleFunc("/healthz", func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte("ok")) + }) + mux.HandleFunc("/events", s.HandleEvents) + return mux +} + +func (s *Server) HandleEvents(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodPost { + http.Error(w, "method not allowed", http.StatusMethodNotAllowed) + return + } + + var ev MeteringEvent + dec := json.NewDecoder(r.Body) + dec.DisallowUnknownFields() + if err := dec.Decode(&ev); err != nil { + http.Error(w, "invalid json: "+err.Error(), http.StatusBadRequest) + return + } + + if ev.RequestID == "" { + http.Error(w, "missing request_id", http.StatusBadRequest) + return + } + if ev.At.IsZero() { + ev.At = time.Now().UTC() + } + + b, err := json.Marshal(ev) + if err != nil { + http.Error(w, "failed to marshal", http.StatusInternalServerError) + return + } + + s.mu.Lock() + defer s.mu.Unlock() + + if s.w != nil { + _, _ = s.w.Write(b) + _, _ = s.w.WriteString("\n") + _ = s.w.Flush() + } else { + log.Printf("EVENT %s", string(b)) + } + + w.WriteHeader(http.StatusAccepted) + _, _ = w.Write([]byte("accepted")) +} diff --git a/collector/internal/collector/server_test.go b/collector/internal/collector/server_test.go new file mode 100644 index 0000000..6def57c --- /dev/null +++ b/collector/internal/collector/server_test.go @@ -0,0 +1,61 @@ +package collector + +import ( + "bytes" + "encoding/json" + "net/http" + "net/http/httptest" + "testing" + "time" + + "github.com/stretchr/testify/require" +) + +func TestHandleEvents_Accepted(t *testing.T) { + s, err := NewServer("") // stdout mode + require.NoError(t, err) + + ev := MeteringEvent{ + RequestID: "req_123", + Tenant: "default", + Provider: "openai", + Model: "gpt-4", + At: time.Now().UTC(), + } + + body, _ := json.Marshal(ev) + req := httptest.NewRequest(http.MethodPost, "/events", bytes.NewReader(body)) + rec := httptest.NewRecorder() + + s.HandleEvents(rec, req) + + require.Equal(t, http.StatusAccepted, rec.Code) + require.Contains(t, rec.Body.String(), "accepted") +} + +func TestHandleEvents_InvalidJSON(t *testing.T) { + s, err := NewServer("") + require.NoError(t, err) + + req := httptest.NewRequest(http.MethodPost, "/events", bytes.NewBufferString("{bad json")) + rec := httptest.NewRecorder() + + s.HandleEvents(rec, req) + + require.Equal(t, http.StatusBadRequest, rec.Code) +} + +func TestHandleEvents_MissingRequestID(t *testing.T) { + s, err := NewServer("") + require.NoError(t, err) + + ev := MeteringEvent{} + body, _ := json.Marshal(ev) + + req := httptest.NewRequest(http.MethodPost, "/events", bytes.NewReader(body)) + rec := httptest.NewRecorder() + + s.HandleEvents(rec, req) + + require.Equal(t, http.StatusBadRequest, rec.Code) +} diff --git a/collector/internal/collector/types.go b/collector/internal/collector/types.go new file mode 100644 index 0000000..58877ac --- /dev/null +++ b/collector/internal/collector/types.go @@ -0,0 +1,18 @@ +package collector + +import "time" + +type MeteringEvent struct { + RequestID string `json:"request_id"` + Tenant string `json:"tenant"` + AppKey string `json:"app_key"` + Provider string `json:"provider"` + Model string `json:"model"` + PromptTokens int `json:"prompt_tokens"` + CompletionTokens int `json:"completion_tokens"` + TotalTokens int `json:"total_tokens"` + LatencyMs int64 `json:"latency_ms"` + StatusCode int `json:"status_code"` + At time.Time `json:"ts"` + Stream bool `json:"stream,omitempty"` +} diff --git a/collector/internal/collector/util.go b/collector/internal/collector/util.go new file mode 100644 index 0000000..fdea059 --- /dev/null +++ b/collector/internal/collector/util.go @@ -0,0 +1,11 @@ +package collector + +import "os" + +func Getenv(k, def string) string { + v := os.Getenv(k) + if v == "" { + return def + } + return v +} diff --git a/collector/main.go b/collector/main.go deleted file mode 100644 index 9093a4d..0000000 --- a/collector/main.go +++ /dev/null @@ -1,120 +0,0 @@ -package main - -import ( - "bufio" - "encoding/json" - "log" - "net/http" - "os" - "sync" - "time" -) - -type MeteringEvent struct { - RequestID string `json:"request_id"` - Tenant string `json:"tenant"` - AppKey string `json:"app_key"` - Provider string `json:"provider"` - Model string `json:"model"` - PromptTokens int `json:"prompt_tokens"` - CompletionTokens int `json:"completion_tokens"` - TotalTokens int `json:"total_tokens"` - LatencyMs int64 `json:"latency_ms"` - StatusCode int `json:"status_code"` - At time.Time `json:"ts"` - Stream bool `json:"stream,omitempty"` -} - -type Server struct { - mu sync.Mutex - file *os.File - w *bufio.Writer -} - -func main() { - // Default: 8081 to match README/k8s manifests - addr := ":" + getenv("PORT", "8081") - outPath := getenv("EVENT_LOG_PATH", "") // e.g. /data/events.ndjson - - s := &Server{} - if outPath != "" { - f, err := os.OpenFile(outPath, os.O_CREATE|os.O_APPEND|os.O_WRONLY, 0644) - if err != nil { - log.Fatalf("open log file: %v", err) - } - s.file = f - s.w = bufio.NewWriterSize(f, 1<<20) - defer func() { - s.mu.Lock() - defer s.mu.Unlock() - _ = s.w.Flush() - _ = s.file.Close() - }() - log.Printf("collector: writing events to %s", outPath) - } else { - log.Printf("collector: EVENT_LOG_PATH not set; events will be printed to stdout") - } - - mux := http.NewServeMux() - mux.HandleFunc("/healthz", func(w http.ResponseWriter, _ *http.Request) { - w.WriteHeader(http.StatusOK) - _, _ = w.Write([]byte("ok")) - }) - mux.HandleFunc("/events", s.handleEvents) - - log.Printf("collector listening on %s", addr) - log.Fatal(http.ListenAndServe(addr, mux)) -} - -func (s *Server) handleEvents(w http.ResponseWriter, r *http.Request) { - if r.Method != http.MethodPost { - http.Error(w, "method not allowed", http.StatusMethodNotAllowed) - return - } - - var ev MeteringEvent - dec := json.NewDecoder(r.Body) - dec.DisallowUnknownFields() // strict for safety - if err := dec.Decode(&ev); err != nil { - http.Error(w, "invalid json: "+err.Error(), http.StatusBadRequest) - return - } - - // Minimal validation - if ev.RequestID == "" { - http.Error(w, "missing request_id", http.StatusBadRequest) - return - } - if ev.At.IsZero() { - ev.At = time.Now().UTC() - } - - // Persist: NDJSON append OR stdout - b, err := json.Marshal(ev) - if err != nil { - http.Error(w, "failed to marshal", http.StatusInternalServerError) - return - } - - s.mu.Lock() - defer s.mu.Unlock() - - if s.w != nil { - _, _ = s.w.Write(b) - _, _ = s.w.WriteString("\n") - _ = s.w.Flush() - } else { - log.Printf("EVENT %s", string(b)) - } - - w.WriteHeader(http.StatusAccepted) - _, _ = w.Write([]byte("accepted")) -} - -func getenv(k, def string) string { - v := os.Getenv(k) - if v == "" { - return def - } - return v -} diff --git a/proxy/Dockerfile b/proxy/Dockerfile index 22b8d3d..f5e11e6 100644 --- a/proxy/Dockerfile +++ b/proxy/Dockerfile @@ -1,21 +1,21 @@ -# -------- build stage -------- -FROM golang:1.22-alpine AS builder - -WORKDIR /src -COPY go.mod ./ -RUN go mod download - -COPY . . -RUN CGO_ENABLED=0 GOOS=linux GOARCH=amd64 go build -trimpath -ldflags="-s -w" -o /out/llm-proxy . - -# -------- runtime stage -------- -FROM gcr.io/distroless/static:nonroot - -WORKDIR / -COPY --from=builder /out/llm-proxy /llm-proxy - -# default port in code: 8080 -EXPOSE 8080 -USER nonroot:nonroot - -ENTRYPOINT ["/llm-proxy"] +# -------- build stage -------- +FROM golang:1.22-alpine AS builder + +WORKDIR /src +COPY go.mod go.sum ./ +RUN go mod download + +COPY . . +RUN CGO_ENABLED=0 GOOS=linux GOARCH=amd64 \ + go build -trimpath -ldflags="-s -w" -o /out/llm-proxy ./cmd/proxy + +# -------- runtime stage -------- +FROM gcr.io/distroless/static:nonroot + +WORKDIR / +COPY --from=builder /out/llm-proxy /llm-proxy + +EXPOSE 8080 +USER nonroot:nonroot + +ENTRYPOINT ["/llm-proxy"] diff --git a/proxy/cmd/proxy/main.go b/proxy/cmd/proxy/main.go new file mode 100644 index 0000000..c0d7de0 --- /dev/null +++ b/proxy/cmd/proxy/main.go @@ -0,0 +1,32 @@ +package main + +import ( + "log" + "net/http" + "time" + + proxy "llm-proxy/internal/proxy" +) + +func main() { + cfg, err := proxy.LoadConfig() + if err != nil { + log.Fatalf("config error: %v", err) + } + + s := proxy.NewServer(cfg) + + srv := &http.Server{ + Addr: cfg.ListenAddr, + Handler: s.Mux(), + ReadHeaderTimeout: 5 * time.Second, + ReadTimeout: 30 * time.Second, + WriteTimeout: 0, + IdleTimeout: 90 * time.Second, + } + + log.Printf("llm-proxy listening on %s (upstream=%s collector=%s capture_bytes=%d)", + cfg.ListenAddr, cfg.UpstreamBaseURL, cfg.CollectorURL, cfg.MeteringCaptureBytes) + + log.Fatal(srv.ListenAndServe()) +} diff --git a/proxy/go.mod b/proxy/go.mod index bb6a717..fabd7d4 100644 --- a/proxy/go.mod +++ b/proxy/go.mod @@ -1,3 +1,11 @@ module llm-proxy go 1.22 + +require github.com/stretchr/testify v1.9.0 + +require ( + github.com/davecgh/go-spew v1.1.1 // indirect + github.com/pmezard/go-difflib v1.0.0 // indirect + gopkg.in/yaml.v3 v3.0.1 // indirect +) diff --git a/proxy/go.sum b/proxy/go.sum new file mode 100644 index 0000000..60ce688 --- /dev/null +++ b/proxy/go.sum @@ -0,0 +1,10 @@ +github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= +github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= +github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= +github.com/stretchr/testify v1.9.0 h1:HtqpIVDClZ4nwg75+f6Lvsy/wHu+3BoSGCbBAcpTsTg= +github.com/stretchr/testify v1.9.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY= +gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM= +gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= +gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= +gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= diff --git a/proxy/internal/proxy/capture.go b/proxy/internal/proxy/capture.go new file mode 100644 index 0000000..7f4d1ed --- /dev/null +++ b/proxy/internal/proxy/capture.go @@ -0,0 +1,33 @@ +package proxy + +type limitedCapture struct { + limit int + buf []byte +} + +func NewLimitedCapture(limit int) *limitedCapture { + if limit <= 0 { + return &limitedCapture{limit: 0, buf: nil} + } + return &limitedCapture{limit: limit, buf: make([]byte, 0, Min(limit, 16*1024))} +} + +func (lc *limitedCapture) Write(p []byte) (int, error) { + if lc.limit <= 0 { + return len(p), nil + } + remain := lc.limit - len(lc.buf) + if remain <= 0 { + return len(p), nil + } + if len(p) <= remain { + lc.buf = append(lc.buf, p...) + return len(p), nil + } + lc.buf = append(lc.buf, p[:remain]...) + return len(p), nil +} + +func (lc *limitedCapture) Bytes() []byte { + return lc.buf +} diff --git a/proxy/internal/proxy/capture_test.go b/proxy/internal/proxy/capture_test.go new file mode 100644 index 0000000..22bad90 --- /dev/null +++ b/proxy/internal/proxy/capture_test.go @@ -0,0 +1,22 @@ +package proxy + +import ( + "testing" + + "github.com/stretchr/testify/require" +) + +func TestLimitedCapture_RespectsLimit(t *testing.T) { + lc := NewLimitedCapture(5) + n, err := lc.Write([]byte("hello world")) + require.NoError(t, err) + require.Equal(t, 11, n) + require.Equal(t, []byte("hello"), lc.Bytes()) +} + +func TestLimitedCapture_ZeroLimit(t *testing.T) { + lc := NewLimitedCapture(0) + _, err := lc.Write([]byte("data")) + require.NoError(t, err) + require.Nil(t, lc.Bytes()) +} diff --git a/proxy/internal/proxy/config.go b/proxy/internal/proxy/config.go new file mode 100644 index 0000000..10fc06f --- /dev/null +++ b/proxy/internal/proxy/config.go @@ -0,0 +1,28 @@ +package proxy + +import ( + "errors" + "os" + "time" +) + +func LoadConfig() (Config, error) { + cfg := Config{ + ListenAddr: EnvOr("LISTEN_ADDR", ":8080"), + UpstreamBaseURL: EnvOr("UPSTREAM_OPENAI_BASE_URL", "https://api.openai.com"), + UpstreamAPIKey: os.Getenv("UPSTREAM_OPENAI_API_KEY"), + CollectorURL: EnvOr("COLLECTOR_URL", "http://llm-collector.llm-system.svc.cluster.local:8081/events"), + EventQueueSize: EnvOrInt("EVENT_QUEUE_SIZE", 10000), + EventFlushTimeout: EnvOrDuration("EVENT_FLUSH_TIMEOUT", 2*time.Second), + HTTPClientTimeout: EnvOrDuration("HTTP_CLIENT_TIMEOUT", 120*time.Second), + MeteringCaptureBytes: EnvOrInt("METERING_CAPTURE_BYTES", 256*1024), + } + + if cfg.UpstreamAPIKey == "" { + return cfg, errors.New("UPSTREAM_OPENAI_API_KEY is required") + } + if cfg.MeteringCaptureBytes < 0 { + cfg.MeteringCaptureBytes = 0 + } + return cfg, nil +} diff --git a/proxy/internal/proxy/handler.go b/proxy/internal/proxy/handler.go new file mode 100644 index 0000000..2134cab --- /dev/null +++ b/proxy/internal/proxy/handler.go @@ -0,0 +1,288 @@ +package proxy + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "io" + "log" + "net" + "net/http" + "strings" + "sync/atomic" + "time" +) + +type Server struct { + cfg Config + upstreamClient *http.Client + collectorClient *http.Client + + events chan MeteringEvent + dropped uint64 +} + +func NewServer(cfg Config) *Server { + transport := &http.Transport{ + Proxy: http.ProxyFromEnvironment, + DialContext: (&net.Dialer{ + Timeout: 5 * time.Second, + KeepAlive: 30 * time.Second, + }).DialContext, + MaxIdleConns: 200, + MaxIdleConnsPerHost: 200, + IdleConnTimeout: 90 * time.Second, + TLSHandshakeTimeout: 5 * time.Second, + ExpectContinueTimeout: 1 * time.Second, + ForceAttemptHTTP2: true, + } + + s := &Server{ + cfg: cfg, + upstreamClient: &http.Client{ + Timeout: cfg.HTTPClientTimeout, + Transport: transport, + }, + collectorClient: &http.Client{ + Timeout: 800 * time.Millisecond, + Transport: transport, + }, + events: make(chan MeteringEvent, cfg.EventQueueSize), + } + + go s.backgroundSender() + + return s +} + +func (s *Server) backgroundSender() { + ticker := time.NewTicker(s.cfg.EventFlushTimeout) + defer ticker.Stop() + + for { + select { + case ev := <-s.events: + if err := postEvent(s.collectorClient, s.cfg.CollectorURL, ev); err != nil { + log.Printf("collector post failed (drop): %v", err) + } + case <-ticker.C: + d := atomic.LoadUint64(&s.dropped) + if d > 0 { + log.Printf("metering: dropped_events=%d (queue full or overload)", d) + } + } + } +} + +func (s *Server) Mux() *http.ServeMux { + mux := http.NewServeMux() + + mux.HandleFunc("/healthz", func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte("ok")) + }) + + mux.HandleFunc("/v1/chat/completions", s.handleChatCompletions) + + return mux +} + +func (s *Server) handleChatCompletions(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodPost { + http.Error(w, "method not allowed", http.StatusMethodNotAllowed) + return + } + + start := time.Now() + requestID := NewReqID() + + tenant := FirstNonEmpty( + r.Header.Get("X-LLM-Tenant"), + r.Header.Get("X-Tenant"), + "default", + ) + + appKey := BearerToken(r.Header.Get("Authorization")) + if appKey == "" { + http.Error(w, "missing Authorization bearer token (gateway key)", http.StatusUnauthorized) + return + } + + reqBody, err := io.ReadAll(io.LimitReader(r.Body, 8<<20)) + if err != nil { + http.Error(w, "failed to read body", http.StatusBadRequest) + return + } + _ = r.Body.Close() + + var oreq OpenAIRequest + _ = json.Unmarshal(reqBody, &oreq) + + upURL := strings.TrimRight(s.cfg.UpstreamBaseURL, "/") + "/v1/chat/completions" + upReq, err := http.NewRequestWithContext(r.Context(), http.MethodPost, upURL, bytes.NewReader(reqBody)) + if err != nil { + http.Error(w, "failed to create upstream request", http.StatusInternalServerError) + return + } + upReq.Header.Set("Content-Type", "application/json") + upReq.Header.Set("Authorization", "Bearer "+s.cfg.UpstreamAPIKey) + + if org := r.Header.Get("OpenAI-Organization"); org != "" { + upReq.Header.Set("OpenAI-Organization", org) + } + if beta := r.Header.Get("OpenAI-Beta"); beta != "" { + upReq.Header.Set("OpenAI-Beta", beta) + } + if proj := r.Header.Get("OpenAI-Project"); proj != "" { + upReq.Header.Set("OpenAI-Project", proj) + } + + upResp, err := s.upstreamClient.Do(upReq) + if err != nil { + http.Error(w, "upstream request failed", http.StatusBadGateway) + s.enqueue(MeteringEvent{ + RequestID: requestID, + Tenant: tenant, + AppKey: appKey, + Provider: "openai", + Model: FirstNonEmpty(oreq.Model, "unknown"), + LatencyMs: time.Since(start).Milliseconds(), + StatusCode: 0, + At: time.Now().UTC(), + PromptTokens: 0, + CompletionTokens: 0, + TotalTokens: 0, + }) + return + } + defer upResp.Body.Close() + + for k, vals := range upResp.Header { + if IsHopByHopHeader(k) { + continue + } + for _, v := range vals { + w.Header().Add(k, v) + } + } + w.Header().Set("X-LLM-Request-ID", requestID) + w.WriteHeader(upResp.StatusCode) + + if oreq.Stream { + seenModel, seenUsage, copyErr := StreamSSE(w, upResp.Body) + lat := time.Since(start) + + model := FirstNonEmpty(seenModel, oreq.Model, "unknown") + ev := MeteringEvent{ + RequestID: requestID, + Tenant: tenant, + AppKey: appKey, + Provider: "openai", + Model: model, + LatencyMs: lat.Milliseconds(), + StatusCode: upResp.StatusCode, + At: time.Now().UTC(), + PromptTokens: 0, + CompletionTokens: 0, + TotalTokens: 0, + } + if seenUsage != nil { + ev.PromptTokens = seenUsage.PromptTokens + ev.CompletionTokens = seenUsage.CompletionTokens + ev.TotalTokens = seenUsage.TotalTokens + } + s.enqueue(ev) + + if copyErr != nil { + log.Printf("proxy stream copy error request_id=%s status=%d err=%v", requestID, upResp.StatusCode, copyErr) + } + return + } + + capWriter := NewLimitedCapture(s.cfg.MeteringCaptureBytes) + tee := io.TeeReader(upResp.Body, capWriter) + + var out io.Writer = w + if fl, ok := w.(http.Flusher); ok { + out = &flushWriter{w: w, fl: fl} + } + + _, copyErr := io.Copy(out, tee) + lat := time.Since(start) + + captured := capWriter.Bytes() + var oresp OpenAIResponse + if len(captured) > 0 { + _ = json.Unmarshal(captured, &oresp) + } + + model := FirstNonEmpty(oresp.Model, oreq.Model, "unknown") + + ev := MeteringEvent{ + RequestID: requestID, + Tenant: tenant, + AppKey: appKey, + Provider: "openai", + Model: model, + LatencyMs: lat.Milliseconds(), + StatusCode: upResp.StatusCode, + At: time.Now().UTC(), + PromptTokens: 0, + CompletionTokens: 0, + TotalTokens: 0, + } + if oresp.Usage != nil { + ev.PromptTokens = oresp.Usage.PromptTokens + ev.CompletionTokens = oresp.Usage.CompletionTokens + ev.TotalTokens = oresp.Usage.TotalTokens + } + s.enqueue(ev) + + if copyErr != nil { + log.Printf("proxy copy error request_id=%s status=%d err=%v", requestID, upResp.StatusCode, copyErr) + } +} + +func (s *Server) enqueue(ev MeteringEvent) { + select { + case s.events <- ev: + default: + atomic.AddUint64(&s.dropped, 1) + } +} + +func postEvent(client *http.Client, collectorURL string, ev MeteringEvent) error { + b, err := json.Marshal(ev) + if err != nil { + return err + } + req, err := http.NewRequestWithContext(context.Background(), http.MethodPost, collectorURL, bytes.NewReader(b)) + if err != nil { + return err + } + req.Header.Set("Content-Type", "application/json") + + resp, err := client.Do(req) + if err != nil { + return err + } + defer resp.Body.Close() + + if resp.StatusCode/100 != 2 { + body, _ := io.ReadAll(io.LimitReader(resp.Body, 4<<10)) + return errors.New("collector returned " + resp.Status + " body=" + string(body)) + } + return nil +} + +type flushWriter struct { + w http.ResponseWriter + fl http.Flusher +} + +func (fw *flushWriter) Write(p []byte) (int, error) { + n, err := fw.w.Write(p) + fw.fl.Flush() + return n, err +} diff --git a/proxy/internal/proxy/sse.go b/proxy/internal/proxy/sse.go new file mode 100644 index 0000000..22042fe --- /dev/null +++ b/proxy/internal/proxy/sse.go @@ -0,0 +1,62 @@ +package proxy + +import ( + "bufio" + "bytes" + "encoding/json" + "errors" + "io" + "net/http" +) + +func StreamSSE(w http.ResponseWriter, upstream io.Reader) (string, *Usage, error) { + br := bufio.NewReaderSize(upstream, 32*1024) + + var model string + var usage *Usage + + var fl http.Flusher + if f, ok := w.(http.Flusher); ok { + fl = f + } + + for { + line, err := br.ReadBytes('\n') + + if len(line) > 0 { + if _, werr := w.Write(line); werr != nil { + return model, usage, werr + } + if fl != nil { + fl.Flush() + } + + trim := bytes.TrimSpace(line) + if bytes.HasPrefix(trim, []byte("data:")) { + payload := bytes.TrimSpace(bytes.TrimPrefix(trim, []byte("data:"))) + + if bytes.Equal(payload, []byte("[DONE]")) { + return model, usage, nil + } + if len(payload) > 0 && payload[0] == '{' { + var ch StreamChunk + if jsonErr := json.Unmarshal(payload, &ch); jsonErr == nil { + if ch.Model != "" { + model = ch.Model + } + if ch.Usage != nil { + usage = ch.Usage + } + } + } + } + } + + if err != nil { + if errors.Is(err, io.EOF) { + return model, usage, nil + } + return model, usage, err + } + } +} diff --git a/proxy/internal/proxy/sse_test.go b/proxy/internal/proxy/sse_test.go new file mode 100644 index 0000000..0b8faaa --- /dev/null +++ b/proxy/internal/proxy/sse_test.go @@ -0,0 +1,39 @@ +package proxy + +import ( + "bytes" + "net/http/httptest" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestStreamSSE_MultiChunkWithUsage(t *testing.T) { + input := ` +data: {"id":"1","model":"gpt-4","usage":{"prompt_tokens":5,"completion_tokens":7,"total_tokens":12}} + +data: [DONE] +` + rec := httptest.NewRecorder() + + model, usage, err := StreamSSE(rec, bytes.NewBufferString(input)) + require.NoError(t, err) + require.Equal(t, "gpt-4", model) + require.NotNil(t, usage) + require.Equal(t, 12, usage.TotalTokens) + require.Contains(t, rec.Body.String(), "data:") +} + +func TestStreamSSE_NoUsage(t *testing.T) { + input := ` +data: {"id":"1","model":"gpt-3.5"} + +data: [DONE] +` + rec := httptest.NewRecorder() + + model, usage, err := StreamSSE(rec, bytes.NewBufferString(input)) + require.NoError(t, err) + require.Equal(t, "gpt-3.5", model) + require.Nil(t, usage) +} diff --git a/proxy/internal/proxy/types.go b/proxy/internal/proxy/types.go new file mode 100644 index 0000000..eaf032b --- /dev/null +++ b/proxy/internal/proxy/types.go @@ -0,0 +1,53 @@ +package proxy + +import "time" + +type Config struct { + ListenAddr string + UpstreamBaseURL string + UpstreamAPIKey string + CollectorURL string + EventQueueSize int + EventFlushTimeout time.Duration + HTTPClientTimeout time.Duration + + MeteringCaptureBytes int +} + +type Usage struct { + PromptTokens int `json:"prompt_tokens"` + CompletionTokens int `json:"completion_tokens"` + TotalTokens int `json:"total_tokens"` +} + +type OpenAIResponse struct { + ID string `json:"id"` + Model string `json:"model"` + Usage *Usage `json:"usage"` +} + +type OpenAIRequest struct { + Model string `json:"model"` + Stream bool `json:"stream"` + Messages any `json:"messages"` +} + +type MeteringEvent struct { + RequestID string `json:"request_id"` + Tenant string `json:"tenant"` + AppKey string `json:"app_key"` + Provider string `json:"provider"` + Model string `json:"model"` + PromptTokens int `json:"prompt_tokens"` + CompletionTokens int `json:"completion_tokens"` + TotalTokens int `json:"total_tokens"` + LatencyMs int64 `json:"latency_ms"` + StatusCode int `json:"status_code"` + At time.Time `json:"ts"` +} + +type StreamChunk struct { + ID string `json:"id"` + Model string `json:"model"` + Usage *Usage `json:"usage"` +} diff --git a/proxy/internal/proxy/util.go b/proxy/internal/proxy/util.go new file mode 100644 index 0000000..d8d43c7 --- /dev/null +++ b/proxy/internal/proxy/util.go @@ -0,0 +1,91 @@ +package proxy + +import ( + "crypto/rand" + "encoding/hex" + "os" + "strconv" + "strings" + "time" +) + +func NewReqID() string { + var b [16]byte + if _, err := rand.Read(b[:]); err == nil { + return "req_" + hex.EncodeToString(b[:]) + } + return "req_fallback_" + hex.EncodeToString([]byte(time.Now().UTC().Format(time.RFC3339Nano))) +} + +func EnvOr(key, def string) string { + v := strings.TrimSpace(os.Getenv(key)) + if v == "" { + return def + } + return v +} + +func EnvOrInt(key string, def int) int { + v := strings.TrimSpace(os.Getenv(key)) + if v == "" { + return def + } + n, err := strconv.Atoi(v) + if err != nil { + return def + } + return n +} + +func EnvOrDuration(key string, def time.Duration) time.Duration { + v := strings.TrimSpace(os.Getenv(key)) + if v == "" { + return def + } + d, err := time.ParseDuration(v) + if err != nil { + return def + } + return d +} + +func BearerToken(auth string) string { + auth = strings.TrimSpace(auth) + if auth == "" { + return "" + } + parts := strings.SplitN(auth, " ", 2) + if len(parts) != 2 { + return "" + } + if !strings.EqualFold(parts[0], "Bearer") { + return "" + } + return strings.TrimSpace(parts[1]) +} + +func FirstNonEmpty(vals ...string) string { + for _, v := range vals { + if strings.TrimSpace(v) != "" { + return v + } + } + return "" +} + +func IsHopByHopHeader(k string) bool { + switch strings.ToLower(strings.TrimSpace(k)) { + case "connection", "keep-alive", "proxy-authenticate", "proxy-authorization", + "te", "trailer", "transfer-encoding", "upgrade": + return true + default: + return false + } +} + +func Min(a, b int) int { + if a < b { + return a + } + return b +} diff --git a/proxy/internal/proxy/util_test.go b/proxy/internal/proxy/util_test.go new file mode 100644 index 0000000..57bd8b0 --- /dev/null +++ b/proxy/internal/proxy/util_test.go @@ -0,0 +1,23 @@ +package proxy + +import ( + "testing" + + "github.com/stretchr/testify/require" +) + +func TestBearerToken(t *testing.T) { + require.Equal(t, "abc", BearerToken("Bearer abc")) + require.Equal(t, "", BearerToken("Basic abc")) + require.Equal(t, "", BearerToken("")) +} + +func TestFirstNonEmpty(t *testing.T) { + require.Equal(t, "b", FirstNonEmpty("", "b", "c")) + require.Equal(t, "", FirstNonEmpty("", " ")) +} + +func TestIsHopByHopHeader(t *testing.T) { + require.True(t, IsHopByHopHeader("Connection")) + require.False(t, IsHopByHopHeader("Content-Type")) +} diff --git a/proxy/main.go b/proxy/main.go deleted file mode 100644 index 44d7aeb..0000000 --- a/proxy/main.go +++ /dev/null @@ -1,560 +0,0 @@ -package main - -import ( - "bufio" - "bytes" - "context" - "crypto/rand" - "encoding/hex" - "encoding/json" - "errors" - "io" - "log" - "net" - "net/http" - "os" - "strconv" - "strings" - "sync/atomic" - "time" -) - -type Config struct { - ListenAddr string - UpstreamBaseURL string - UpstreamAPIKey string - CollectorURL string - EventQueueSize int - EventFlushTimeout time.Duration - HTTPClientTimeout time.Duration - - // Best-effort metering parse: capture first N bytes of upstream response - MeteringCaptureBytes int -} - -type Usage struct { - PromptTokens int `json:"prompt_tokens"` - CompletionTokens int `json:"completion_tokens"` - TotalTokens int `json:"total_tokens"` -} - -type OpenAIResponse struct { - ID string `json:"id"` - Model string `json:"model"` - Usage *Usage `json:"usage"` -} - -type OpenAIRequest struct { - Model string `json:"model"` - Stream bool `json:"stream"` - Messages any `json:"messages"` // MVP -} - -type MeteringEvent struct { - RequestID string `json:"request_id"` - Tenant string `json:"tenant"` - AppKey string `json:"app_key"` - Provider string `json:"provider"` - Model string `json:"model"` - PromptTokens int `json:"prompt_tokens"` - CompletionTokens int `json:"completion_tokens"` - TotalTokens int `json:"total_tokens"` - LatencyMs int64 `json:"latency_ms"` - StatusCode int `json:"status_code"` - At time.Time `json:"ts"` -} - -type StreamChunk struct { - ID string `json:"id"` - Model string `json:"model"` - Usage *Usage `json:"usage"` -} - -func loadConfig() (Config, error) { - cfg := Config{ - ListenAddr: envOr("LISTEN_ADDR", ":8080"), - UpstreamBaseURL: envOr("UPSTREAM_OPENAI_BASE_URL", "https://api.openai.com"), - UpstreamAPIKey: os.Getenv("UPSTREAM_OPENAI_API_KEY"), - CollectorURL: envOr("COLLECTOR_URL", "http://llm-collector.llm-system.svc.cluster.local:8081/events"), - EventQueueSize: envOrInt("EVENT_QUEUE_SIZE", 10000), - EventFlushTimeout: envOrDuration("EVENT_FLUSH_TIMEOUT", 2*time.Second), - HTTPClientTimeout: envOrDuration("HTTP_CLIENT_TIMEOUT", 120*time.Second), - MeteringCaptureBytes: envOrInt("METERING_CAPTURE_BYTES", 256*1024), // 256KB - } - - if cfg.UpstreamAPIKey == "" { - return cfg, errors.New("UPSTREAM_OPENAI_API_KEY is required") - } - if cfg.MeteringCaptureBytes < 0 { - cfg.MeteringCaptureBytes = 0 - } - return cfg, nil -} - -func streamSSE(w http.ResponseWriter, upstream io.Reader) (string, *Usage, error) { - br := bufio.NewReaderSize(upstream, 32*1024) - - var model string - var usage *Usage - - var fl http.Flusher - if f, ok := w.(http.Flusher); ok { - fl = f - } - - for { - line, err := br.ReadBytes('\n') // SSE is line-oriented - - if len(line) > 0 { - // 1) forward raw bytes as-is - if _, werr := w.Write(line); werr != nil { - return model, usage, werr // client disconnected etc. - } - if fl != nil { - fl.Flush() - } - - // 2) best-effort parse: data: ... - trim := bytes.TrimSpace(line) - if bytes.HasPrefix(trim, []byte("data:")) { - payload := bytes.TrimSpace(bytes.TrimPrefix(trim, []byte("data:"))) - - // stop marker - if bytes.Equal(payload, []byte("[DONE]")) { - return model, usage, nil - // done, but still continue reading until EOF just in case - } else if len(payload) > 0 && payload[0] == '{' { - var ch StreamChunk - if jsonErr := json.Unmarshal(payload, &ch); jsonErr == nil { - if ch.Model != "" { - model = ch.Model - } - // usage only appears if client enabled stream_options.include_usage - if ch.Usage != nil { - usage = ch.Usage - } - } - } - } - } - - if err != nil { - if errors.Is(err, io.EOF) { - return model, usage, nil - } - return model, usage, err - } - } -} - -func main() { - cfg, err := loadConfig() - if err != nil { - log.Fatalf("config error: %v", err) - } - - // Upstream HTTP client (keep-alive) - transport := &http.Transport{ - Proxy: http.ProxyFromEnvironment, - DialContext: (&net.Dialer{ - Timeout: 5 * time.Second, - KeepAlive: 30 * time.Second, - }).DialContext, - MaxIdleConns: 200, - MaxIdleConnsPerHost: 200, - IdleConnTimeout: 90 * time.Second, - TLSHandshakeTimeout: 5 * time.Second, - ExpectContinueTimeout: 1 * time.Second, - ForceAttemptHTTP2: true, - } - upstreamClient := &http.Client{ - Timeout: cfg.HTTPClientTimeout, - Transport: transport, - } - - // Collector client: short timeout to avoid hanging background sender - collectorClient := &http.Client{ - Timeout: 800 * time.Millisecond, - Transport: transport, - } - - events := make(chan MeteringEvent, cfg.EventQueueSize) - var dropped uint64 - - // Background sender (MVP: one-by-one; later: batch+gzip/protobuf) - go func() { - ticker := time.NewTicker(cfg.EventFlushTimeout) - defer ticker.Stop() - - for { - select { - case ev := <-events: - if err := postEvent(collectorClient, cfg.CollectorURL, ev); err != nil { - // fail-open: drop on error - log.Printf("collector post failed (drop): %v", err) - } - case <-ticker.C: - d := atomic.LoadUint64(&dropped) - if d > 0 { - log.Printf("metering: dropped_events=%d (queue full or overload)", d) - } - } - } - }() - - mux := http.NewServeMux() - mux.HandleFunc("/healthz", func(w http.ResponseWriter, _ *http.Request) { - w.WriteHeader(http.StatusOK) - _, _ = w.Write([]byte("ok")) - }) - - // MVP endpoint: Chat Completions - mux.HandleFunc("/v1/chat/completions", func(w http.ResponseWriter, r *http.Request) { - if r.Method != http.MethodPost { - http.Error(w, "method not allowed", http.StatusMethodNotAllowed) - return - } - - start := time.Now() - requestID := newReqID() - - tenant := firstNonEmpty( - r.Header.Get("X-LLM-Tenant"), - r.Header.Get("X-Tenant"), - "default", - ) - - appKey := bearerToken(r.Header.Get("Authorization")) - if appKey == "" { - http.Error(w, "missing Authorization bearer token (gateway key)", http.StatusUnauthorized) - return - } - - // Read request body (kept for MVP). Next iteration can stream request too. - reqBody, err := io.ReadAll(io.LimitReader(r.Body, 8<<20)) - if err != nil { - http.Error(w, "failed to read body", http.StatusBadRequest) - return - } - _ = r.Body.Close() - - // Parse minimal fields from request (model/stream) - var oreq OpenAIRequest - _ = json.Unmarshal(reqBody, &oreq) - - // Build upstream request - upURL := strings.TrimRight(cfg.UpstreamBaseURL, "/") + "/v1/chat/completions" - upReq, err := http.NewRequestWithContext(r.Context(), http.MethodPost, upURL, bytes.NewReader(reqBody)) - if err != nil { - http.Error(w, "failed to create upstream request", http.StatusInternalServerError) - return - } - upReq.Header.Set("Content-Type", "application/json") - upReq.Header.Set("Authorization", "Bearer "+cfg.UpstreamAPIKey) - - // Forward some optional headers if needed - if org := r.Header.Get("OpenAI-Organization"); org != "" { - upReq.Header.Set("OpenAI-Organization", org) - } - if beta := r.Header.Get("OpenAI-Beta"); beta != "" { - upReq.Header.Set("OpenAI-Beta", beta) - } - if proj := r.Header.Get("OpenAI-Project"); proj != "" { - upReq.Header.Set("OpenAI-Project", proj) - } - - upResp, err := upstreamClient.Do(upReq) - if err != nil { - http.Error(w, "upstream request failed", http.StatusBadGateway) - - // Emit metering event even on upstream failure (status_code=0) - ev := MeteringEvent{ - RequestID: requestID, - Tenant: tenant, - AppKey: appKey, - Provider: "openai", - Model: firstNonEmpty(oreq.Model, "unknown"), - PromptTokens: 0, - CompletionTokens: 0, - TotalTokens: 0, - LatencyMs: time.Since(start).Milliseconds(), - StatusCode: 0, - At: time.Now().UTC(), - } - enqueueEvent(events, &dropped, ev) - return - } - defer upResp.Body.Close() - - // Copy headers (filter hop-by-hop) - for k, vals := range upResp.Header { - if isHopByHopHeader(k) { - continue - } - for _, v := range vals { - w.Header().Add(k, v) - } - } - // Add our own trace id to help debugging - w.Header().Set("X-LLM-Request-ID", requestID) - - // Write status code immediately - w.WriteHeader(upResp.StatusCode) - - if oreq.Stream { - seenModel, seenUsage, copyErr := streamSSE(w, upResp.Body) - lat := time.Since(start) - - model := firstNonEmpty(seenModel, oreq.Model, "unknown") - - ev := MeteringEvent{ - RequestID: requestID, - Tenant: tenant, - AppKey: appKey, - Provider: "openai", - Model: model, - PromptTokens: 0, - CompletionTokens: 0, - TotalTokens: 0, - LatencyMs: lat.Milliseconds(), - StatusCode: upResp.StatusCode, - At: time.Now().UTC(), - } - if seenUsage != nil { - ev.PromptTokens = seenUsage.PromptTokens - ev.CompletionTokens = seenUsage.CompletionTokens - ev.TotalTokens = seenUsage.TotalTokens - } - - enqueueEvent(events, &dropped, ev) - - if copyErr != nil { - log.Printf("proxy stream copy error request_id=%s status=%d err=%v", requestID, upResp.StatusCode, copyErr) - } - return - } - - // We want: stream upstream -> client WITHOUT buffering. - // Additionally: capture first N bytes for best-effort usage parsing. - capWriter := newLimitedCapture(cfg.MeteringCaptureBytes) - tee := io.TeeReader(upResp.Body, capWriter) - - // Ensure streaming to client flushes when possible - var out io.Writer = w - if fl, ok := w.(http.Flusher); ok { - out = &flushWriter{w: w, fl: fl} - } - - _, copyErr := io.Copy(out, tee) - lat := time.Since(start) - - // Best-effort parse usage from captured bytes (may fail if truncated) - captured := capWriter.Bytes() - var oresp OpenAIResponse - if len(captured) > 0 { - _ = json.Unmarshal(captured, &oresp) - } - - model := firstNonEmpty(oresp.Model, oreq.Model, "unknown") - - ev := MeteringEvent{ - RequestID: requestID, - Tenant: tenant, - AppKey: appKey, - Provider: "openai", - Model: model, - PromptTokens: 0, - CompletionTokens: 0, - TotalTokens: 0, - LatencyMs: lat.Milliseconds(), - StatusCode: upResp.StatusCode, - At: time.Now().UTC(), - } - if oresp.Usage != nil { - ev.PromptTokens = oresp.Usage.PromptTokens - ev.CompletionTokens = oresp.Usage.CompletionTokens - ev.TotalTokens = oresp.Usage.TotalTokens - } - - enqueueEvent(events, &dropped, ev) - - // If copy failed, log it (client might have disconnected) - if copyErr != nil { - log.Printf("proxy copy error request_id=%s status=%d err=%v", requestID, upResp.StatusCode, copyErr) - } - }) - - srv := &http.Server{ - Addr: cfg.ListenAddr, - Handler: mux, - ReadHeaderTimeout: 5 * time.Second, - ReadTimeout: 30 * time.Second, - WriteTimeout: 0, // response can be long; keep 0 for now - IdleTimeout: 90 * time.Second, - } - - log.Printf("llm-proxy listening on %s (upstream=%s collector=%s capture_bytes=%d)", - cfg.ListenAddr, cfg.UpstreamBaseURL, cfg.CollectorURL, cfg.MeteringCaptureBytes) - - log.Fatal(srv.ListenAndServe()) -} - -func enqueueEvent(ch chan MeteringEvent, dropped *uint64, ev MeteringEvent) { - select { - case ch <- ev: - default: - atomic.AddUint64(dropped, 1) - } -} - -func postEvent(client *http.Client, collectorURL string, ev MeteringEvent) error { - b, err := json.Marshal(ev) - if err != nil { - return err - } - req, err := http.NewRequestWithContext(context.Background(), http.MethodPost, collectorURL, bytes.NewReader(b)) - if err != nil { - return err - } - req.Header.Set("Content-Type", "application/json") - - resp, err := client.Do(req) - if err != nil { - return err - } - defer resp.Body.Close() - - if resp.StatusCode/100 != 2 { - body, _ := io.ReadAll(io.LimitReader(resp.Body, 4<<10)) - return errors.New("collector returned " + resp.Status + " body=" + string(body)) - } - return nil -} - -func newReqID() string { - var b [16]byte - if _, err := rand.Read(b[:]); err == nil { - return "req_" + hex.EncodeToString(b[:]) - } - return "req_fallback_" + hex.EncodeToString([]byte(time.Now().UTC().Format(time.RFC3339Nano))) -} - -func envOr(key, def string) string { - v := strings.TrimSpace(os.Getenv(key)) - if v == "" { - return def - } - return v -} - -func envOrInt(key string, def int) int { - v := strings.TrimSpace(os.Getenv(key)) - if v == "" { - return def - } - n, err := strconv.Atoi(v) - if err != nil { - return def - } - return n -} - -func envOrDuration(key string, def time.Duration) time.Duration { - v := strings.TrimSpace(os.Getenv(key)) - if v == "" { - return def - } - d, err := time.ParseDuration(v) - if err != nil { - return def - } - return d -} - -func bearerToken(auth string) string { - auth = strings.TrimSpace(auth) - if auth == "" { - return "" - } - parts := strings.SplitN(auth, " ", 2) - if len(parts) != 2 { - return "" - } - if !strings.EqualFold(parts[0], "Bearer") { - return "" - } - return strings.TrimSpace(parts[1]) -} - -func firstNonEmpty(vals ...string) string { - for _, v := range vals { - if strings.TrimSpace(v) != "" { - return v - } - } - return "" -} - -func isHopByHopHeader(k string) bool { - switch strings.ToLower(strings.TrimSpace(k)) { - case "connection", "keep-alive", "proxy-authenticate", "proxy-authorization", - "te", "trailer", "transfer-encoding", "upgrade": - return true - default: - return false - } -} - -// flushWriter flushes after each Write to reduce buffering and improve TTFB. -// This does not change the response format; clients still see a single JSON body. -type flushWriter struct { - w http.ResponseWriter - fl http.Flusher -} - -func (fw *flushWriter) Write(p []byte) (int, error) { - n, err := fw.w.Write(p) - fw.fl.Flush() - return n, err -} - -// limitedCapture collects up to N bytes written to it, ignoring the rest. -// Used to "tee" upstream response for best-effort JSON usage parsing without buffering everything. -type limitedCapture struct { - limit int - buf []byte -} - -func newLimitedCapture(limit int) *limitedCapture { - if limit <= 0 { - return &limitedCapture{limit: 0, buf: nil} - } - return &limitedCapture{limit: limit, buf: make([]byte, 0, min(limit, 16*1024))} -} - -func (lc *limitedCapture) Write(p []byte) (int, error) { - if lc.limit <= 0 { - return len(p), nil - } - remain := lc.limit - len(lc.buf) - if remain <= 0 { - return len(p), nil - } - if len(p) <= remain { - lc.buf = append(lc.buf, p...) - return len(p), nil - } - lc.buf = append(lc.buf, p[:remain]...) - return len(p), nil -} - -func (lc *limitedCapture) Bytes() []byte { - return lc.buf -} - -func min(a, b int) int { - if a < b { - return a - } - return b -}