From dcfcfb7e48aec20e8c3d860d7b98e8a29cceecd2 Mon Sep 17 00:00:00 2001 From: Lukasz Raczylo Date: Thu, 20 Aug 2026 10:45:12 +0100 Subject: [PATCH] feat: unified chat and embeddings endpoints with limits and accounting Implements Task 12: POST /v1/chat/completions and POST /v1/embeddings through a shared pipeline (routes_unified.go: handleChat, handleEmbeddings, runUnified) wired into ServeHTTP's dispatch in llmgateway.go. Pipeline: decode (10MiB cap) -> resolve model -> enforce user/group limits -> call adapter -> account usage, translating every failure mode into its HTTP envelope per the task brief's carried rulings: - ALIAS ECHO (ruling a): the client-requested model id is threaded to the adapter via the internal gatewayAliasKey convention and echoed back in anthropic/gemini's translated response model field; every adapter strips the key before building its upstream wire body so it never leaks to a real provider. - Limit violations (ruling b): storeDown -> 503 server_error; otherwise 429 rate_limit_error with Retry-After when the violation names a window. - Stream-safe error handling (ruling c): statusTrackingWriter.wroteHeader gates whether an adapter error still gets an envelope; providerHTTPError wraps into {error:{message,type,code,upstream}}; translateError maps to 400/501; context.Canceled logs quietly; anything else is a 502. - Usage/accounting (ruling d): zero non-streaming usage gets a ceil(len(body)/4) prompt-token estimate; zero streaming usage logs once and accounts the request only. unifiedCostMicros extends costMicros' single-id lookup to canonical-then-bare without changing its signature. - Scopes (ruling e): user scope only when u.limits is set, group scope only when grp.limits is set, user first. routes_unified_test.go drives every case end-to-end through Gateway.ServeHTTP against httptest fake upstreams (openai- and anthropic-shaped): happy path with counters asserted via the limiter's own store, incremental stream flushing, 401/403/404/429/503, embeddings, alias echo through a provider-prefixed anthropic request, wrapped upstream error envelopes, and a mid-stream drop that must not append a second envelope to an already-committed response. Gates: gofmt, go vet, fieldalignment, golangci-lint (incl. gosec) all clean; go test ./... -race -count=1 passes. --- llmgateway.go | 14 ++ provider_anthropic.go | 8 + provider_gemini.go | 23 +- provider_openai.go | 9 + routes_unified.go | 242 ++++++++++++++++++ routes_unified_test.go | 558 +++++++++++++++++++++++++++++++++++++++++ 6 files changed, 851 insertions(+), 3 deletions(-) create mode 100644 routes_unified_test.go diff --git a/llmgateway.go b/llmgateway.go index 1c17fdb..aab5c9d 100644 --- a/llmgateway.go +++ b/llmgateway.go @@ -254,6 +254,20 @@ func (g *Gateway) ServeHTTP(w http.ResponseWriter, r *http.Request) { return } + if r.Method == http.MethodPost && (r.URL.Path == "/v1/chat/completions" || r.URL.Path == "/v1/embeddings") { + u, grp, ok := g.auth.identify(r) + if !ok { + writeOAIError(sw, http.StatusUnauthorized, "authentication_error", "invalid or missing API key") + return + } + if r.URL.Path == "/v1/chat/completions" { + g.handleChat(sw, r, u, grp) + } else { + g.handleEmbeddings(sw, r, u, grp) + } + return + } + if g.cfg.PassthroughUnknown { g.next.ServeHTTP(sw, r) return diff --git a/provider_anthropic.go b/provider_anthropic.go index bff0058..3d24b50 100644 --- a/provider_anthropic.go +++ b/provider_anthropic.go @@ -81,6 +81,14 @@ func (a *anthropicAdapter) requestHeaders(hasBody bool) http.Header { // request is made. func (a *anthropicAdapter) chatCompletion(ctx context.Context, w http.ResponseWriter, req map[string]any) (usage, error) { gatewayModel, _ := req["model"].(string) + // gatewayAliasKey (ruling a, ALIAS ECHO), when present, is the exact id + // the client requested — echoed into the translated response's "model" + // field below instead of the bare upstream id, then deleted so it + // cannot leak into the upstream request body. + if alias, ok := req[gatewayAliasKey].(string); ok && alias != "" { + gatewayModel = alias + } + delete(req, gatewayAliasKey) streaming, _ := req["stream"].(bool) body, err := anthropicRequestFromOpenAI(req) diff --git a/provider_gemini.go b/provider_gemini.go index 818786f..dfe414a 100644 --- a/provider_gemini.go +++ b/provider_gemini.go @@ -84,6 +84,16 @@ func geminiModelFromRequest(req map[string]any) string { // request is made. func (a *geminiAdapter) chatCompletion(ctx context.Context, w http.ResponseWriter, req map[string]any) (usage, error) { model := geminiModelFromRequest(req) + // responseModel (ruling a, ALIAS ECHO): gatewayAliasKey, when present, + // is the exact id the client requested — echoed into the translated + // response's "model" field, distinct from model (the bare upstream id), + // which must keep driving the URL below unchanged. Deleted from req so + // it cannot leak into the upstream request body. + responseModel := model + if alias, ok := req[gatewayAliasKey].(string); ok && alias != "" { + responseModel = alias + } + delete(req, gatewayAliasKey) streaming, _ := req["stream"].(bool) body, err := geminiRequestFromOpenAI(req) @@ -117,11 +127,11 @@ func (a *geminiAdapter) chatCompletion(ctx context.Context, w http.ResponseWrite } if streaming && strings.Contains(strings.ToLower(resp.Header.Get("Content-Type")), "text/event-stream") { - return a.forwardStream(w, resp.Body, model) + return a.forwardStream(w, resp.Body, responseModel) } // Either a non-streaming request, or a streaming one whose upstream // ignored stream=true and answered with a normal JSON body. - return a.forwardJSON(w, resp, model) + return a.forwardJSON(w, resp, responseModel) } // forwardJSON reads resp's full body (capped at maxResponseBytes), @@ -189,6 +199,13 @@ func (a *geminiAdapter) forwardStream(w http.ResponseWriter, body io.Reader, gat // Embeddings never stream. func (a *geminiAdapter) embeddings(ctx context.Context, w http.ResponseWriter, req map[string]any) (usage, error) { model := geminiModelFromRequest(req) + // responseModel (ruling a, ALIAS ECHO): see chatCompletion's identical + // pattern above. + responseModel := model + if alias, ok := req[gatewayAliasKey].(string); ok && alias != "" { + responseModel = alias + } + delete(req, gatewayAliasKey) texts, single, err := geminiEmbeddingInputTexts(req["input"]) if err != nil { @@ -223,7 +240,7 @@ func (a *geminiAdapter) embeddings(ctx context.Context, w http.ResponseWriter, r return usage{}, fmt.Errorf("%w: read response body: %w", errUpstream, err) } - out, u, err := openAIEmbeddingResponseFromGemini(raw, single, model, estimatedEmbeddingTokens(texts)) + out, u, err := openAIEmbeddingResponseFromGemini(raw, single, responseModel, estimatedEmbeddingTokens(texts)) if err != nil { return usage{}, err } diff --git a/provider_openai.go b/provider_openai.go index 0b9fb07..b326968 100644 --- a/provider_openai.go +++ b/provider_openai.go @@ -92,6 +92,14 @@ type chatStreamChunk struct { // resulting usage-only chunk from a client that never asked for it while // still capturing its usage for accounting. func (a *openaiAdapter) chatCompletion(ctx context.Context, w http.ResponseWriter, req map[string]any) (usage, error) { + // gatewayAliasKey must never reach the real provider: an openai-type + // adapter forwards req verbatim as the upstream wire body, so a + // leftover alias entry would arrive as an unrecognized request field. + // This adapter's own response is a verbatim passthrough of whatever + // the upstream returns, so — unlike anthropic/gemini — there is no + // alias to echo back into it (documented in gatewayAliasKey's doc + // comment, routes_unified.go). + delete(req, gatewayAliasKey) streaming, _ := req["stream"].(bool) clientAskedUsage := false if streaming { @@ -135,6 +143,7 @@ func (a *openaiAdapter) chatCompletion(ctx context.Context, w http.ResponseWrite // stream, so it is forwardJSON end to end, same as chatCompletion's // non-streaming path. func (a *openaiAdapter) embeddings(ctx context.Context, w http.ResponseWriter, req map[string]any) (usage, error) { + delete(req, gatewayAliasKey) // see chatCompletion's identical delete for why resp, err := upstreamJSON(ctx, a.client, http.MethodPost, a.baseURL+"/v1/embeddings", a.requestHeaders(true), req) if err != nil { return usage{}, err diff --git a/routes_unified.go b/routes_unified.go index f08aabe..7011a6d 100644 --- a/routes_unified.go +++ b/routes_unified.go @@ -1,10 +1,41 @@ package traefikllmgateway import ( + "context" "encoding/json" + "errors" + "io" + "math" "net/http" + "strconv" ) +// maxRequestBytes caps a client request body the unified routes decode +// into a map[string]any: 10MiB, generous for a chat or embeddings request +// (including an inline base64 image) while bounding memory against an +// oversized or malicious body. +const maxRequestBytes = 10 << 20 + +// gatewayAliasKey is the request map's internal-convention key carrying +// the client-requested model id through to an adapter (ruling a, ALIAS +// ECHO), distinct from "model" itself, which runUnified rewrites to the +// resolved upstream model id before calling the adapter — an operator's +// model alias must never reach the upstream API, but the client that used +// it must still see it echoed back in the response. +// +// Every adapter deletes this key before marshaling its upstream wire +// body: an openai-type adapter forwards req verbatim, so a leftover +// gatewayAliasKey entry would otherwise reach the real provider as an +// unrecognized request field. anthropic and gemini's adapters, which +// build a brand new response envelope rather than forwarding one +// verbatim, read it (falling back to the upstream model id when absent) +// to set that envelope's "model" field to the client's original id. An +// openai-type adapter's own response is a verbatim passthrough of +// whatever the upstream returned and cannot be rewritten this way — its +// response body carries the upstream's model id, not the alias. This is +// a documented, accepted asymmetry, not a bug. +const gatewayAliasKey = "__alias" + // handleModels implements GET /v1/models: an authenticated caller gets back // an OpenAI-compatible {"object":"list","data":[...]} envelope of the // models their group can see, via modelRegistry.listFor. @@ -21,3 +52,214 @@ func (g *Gateway) handleModels(w http.ResponseWriter, r *http.Request) { "data": g.registry.listFor(grp), }) } + +// adapterCall abstracts providerAdapter.chatCompletion and .embeddings: +// handleChat and handleEmbeddings share every other step of the pipeline +// below (decode, resolve, enforce limits, account usage, translate +// errors) and differ only in which of these two methods they invoke. +type adapterCall func(a providerAdapter, ctx context.Context, w http.ResponseWriter, req map[string]any) (usage, error) + +// handleChat implements POST /v1/chat/completions. +func (g *Gateway) handleChat(w http.ResponseWriter, r *http.Request, u *user, grp *group) { + g.runUnified(w, r, u, grp, func(a providerAdapter, ctx context.Context, w http.ResponseWriter, req map[string]any) (usage, error) { + return a.chatCompletion(ctx, w, req) + }) +} + +// handleEmbeddings implements POST /v1/embeddings. +func (g *Gateway) handleEmbeddings(w http.ResponseWriter, r *http.Request, u *user, grp *group) { + g.runUnified(w, r, u, grp, func(a providerAdapter, ctx context.Context, w http.ResponseWriter, req map[string]any) (usage, error) { + return a.embeddings(ctx, w, req) + }) +} + +// runUnified is the shared chat/embeddings pipeline: decode the request, +// resolve its model, enforce per-user and per-group limits, invoke the +// adapter via call, then account the resulting usage. w is wrapped in its +// own statusTrackingWriter so a mid-stream adapter error (headers already +// sent) can be told apart from one that failed before any write. +func (g *Gateway) runUnified(w http.ResponseWriter, r *http.Request, u *user, grp *group, call adapterCall) { + sw := &statusTrackingWriter{ResponseWriter: w} + + body, err := io.ReadAll(io.LimitReader(r.Body, maxRequestBytes)) + if err != nil { + writeOAIError(sw, http.StatusBadRequest, "invalid_request_error", "cannot read request body") + return + } + + var req map[string]any + if err = json.Unmarshal(body, &req); err != nil { + writeOAIError(sw, http.StatusBadRequest, "invalid_request_error", "invalid JSON body") + return + } + + requestedModel, _ := req["model"].(string) + if requestedModel == "" { + writeOAIError(sw, http.StatusBadRequest, "invalid_request_error", "model is required") + return + } + streaming, _ := req["stream"].(bool) + + adapter, upstreamModel, canonical, err := g.registry.resolve(requestedModel, grp) + if err != nil { + writeModelResolveError(sw, err) + return + } + + scopes := buildLimitScopes(u, grp) + if violation := g.limiter.checkAndCount(scopes); violation != nil { + writeLimitViolation(sw, violation) + return + } + + req["model"] = upstreamModel + req[gatewayAliasKey] = requestedModel + + result, callErr := call(adapter, r.Context(), sw, req) + if callErr != nil { + g.handleAdapterError(sw, callErr, adapter.name()) + return + } + + if result.total() == 0 { + if streaming { + g.logf("unified route: zero usage reported for a streaming response from model %q; accounting the request only", canonical) + } else { + result.prompt = int64(math.Ceil(float64(len(body)) / 4)) + result.estimated = true + } + } + + g.limiter.account(scopes, result, unifiedCostMicros(canonical, upstreamModel, result, g.cfg.Pricing)) +} + +// buildLimitScopes returns the limitScope slice runUnified passes to the +// limiter: a user scope only when u has its own limits configured, then a +// group scope only when grp does (ruling e) — either, both, or neither may +// apply to a given request. The user scope is listed first, so +// checkAndCount reports a user's own violation ahead of their group's when +// both are breached by the same request. +func buildLimitScopes(u *user, grp *group) []limitScope { + scopes := make([]limitScope, 0, 2) + if u.limits != nil { + scopes = append(scopes, limitScope{limits: u.limits, kind: "user", id: u.name}) + } + if grp.limits != nil { + scopes = append(scopes, limitScope{limits: grp.limits, kind: "group", id: grp.name}) + } + return scopes +} + +// unifiedCostMicros resolves the price to charge one request's usage +// against: canonical ("provider/model") first, falling back to bare — the +// upstream model id alone — only when canonical has no configured price at +// all, in either overrides or the built-in table. This extends +// costMicros' single-id lookup to the two-id order ruling (d) requires, +// without changing costMicros' own signature: lookupPricing is used purely +// as a side-effect-free existence probe on canonical, and the real +// computation (and its unknown-model warning, when neither id has a price) +// runs through the one costMicros call that is actually charged. +func unifiedCostMicros(canonical, bare string, u usage, overrides map[string]*ModelPricing) int64 { + if _, ok := lookupPricing(canonical, overrides); ok { + return costMicros(canonical, u, overrides) + } + return costMicros(bare, u, overrides) +} + +// writeModelResolveError maps a modelRegistry.resolve error to its HTTP +// envelope: errModelUnknown is a 404 (no configured provider knows this +// model), errModelDenied is a 403 (a real model the caller's group cannot +// use). Any other error is a defensive 500 — resolve's own contract +// promises only these two sentinels, so reaching this branch would be a +// programming error, not a client mistake. +func writeModelResolveError(w http.ResponseWriter, err error) { + switch { + case errors.Is(err, errModelUnknown): + writeOAIError(w, http.StatusNotFound, "invalid_request_error", "unknown model") + case errors.Is(err, errModelDenied): + writeOAIError(w, http.StatusForbidden, "invalid_request_error", "model access denied") + default: + writeOAIError(w, http.StatusInternalServerError, "server_error", "internal error") + } +} + +// writeLimitViolation maps a limiter violation to its HTTP envelope +// (ruling b): a storeDown violation — the configured limit store was +// unreachable and failOpen is false — is a 503 server_error, since it is +// not an actual limit breach; any other violation is a 429 rate_limit_error +// carrying v's own message, plus a Retry-After header when the violation +// names a meaningful retry window. +func writeLimitViolation(w http.ResponseWriter, v *limitViolation) { + if v.storeDown { + writeOAIError(w, http.StatusServiceUnavailable, "server_error", v.message) + return + } + if v.retryAfter > 0 { + w.Header().Set("Retry-After", strconv.Itoa(v.retryAfter)) + } + writeOAIError(w, http.StatusTooManyRequests, "rate_limit_error", v.message) +} + +// handleAdapterError translates an adapter's returned error into a +// response, or into nothing at all, per ruling c. sw.wroteHeader is +// checked first, ahead of every other case: an error surfacing after the +// adapter already started writing a response (a mid-stream connection +// drop) must never get a second, conflicting envelope appended, whatever +// kind of error it is. +func (g *Gateway) handleAdapterError(sw *statusTrackingWriter, err error, providerName string) { + if sw.wroteHeader { + g.errorf("unified route: adapter error after response started (provider %q): %v", providerName, err) + return + } + + var perr *providerHTTPError + if errors.As(err, &perr) { + writeProviderUpstreamError(sw, providerName, perr) + return + } + + var terr *translateError + if errors.As(err, &terr) { + status := http.StatusBadRequest + if terr.notSupported { + status = http.StatusNotImplemented + } + writeOAIError(sw, status, "invalid_request_error", terr.msg) + return + } + + if errors.Is(err, context.Canceled) { + g.logf("unified route: client canceled request to provider %q: %v", providerName, err) + return + } + + g.errorf("unified route: upstream connection error (provider %q): %v", providerName, err) + writeOAIError(sw, http.StatusBadGateway, "server_error", "upstream connection error") +} + +// writeProviderUpstreamError passes perr's status through to the client, +// wrapped in the gateway's own envelope shape rather than perr.body +// verbatim (unified-route callers get this wrapped form; the openai-type +// adapter's own passthrough forwards a non-2xx body unwrapped, since that +// path never reaches this function). perr.body is embedded under +// error.upstream, decoded to a JSON value when it parses as one and left +// as a raw string otherwise, so a caller sees exactly what the provider +// said either way. +func writeProviderUpstreamError(w http.ResponseWriter, providerName string, perr *providerHTTPError) { + var upstream any = string(perr.body) + var parsed any + if json.Unmarshal(perr.body, &parsed) == nil { + upstream = parsed + } + + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(perr.status) + _ = json.NewEncoder(w).Encode(map[string]any{ // headers already committed; nothing useful to do on encode failure + "error": map[string]any{ + "message": providerName + " upstream error", + "type": "upstream_error", + "code": strconv.Itoa(perr.status), + "upstream": upstream, + }, + }) +} diff --git a/routes_unified_test.go b/routes_unified_test.go new file mode 100644 index 0000000..c1b4d98 --- /dev/null +++ b/routes_unified_test.go @@ -0,0 +1,558 @@ +package traefikllmgateway + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "io" + "net/http" + "net/http/httptest" + "os" + "strings" + "testing" + "time" +) + +// newUnifiedRequest builds a POST request against the unified routes, +// JSON-encoding body and attaching apiKey as a Bearer token when non-empty. +func newUnifiedRequest(t *testing.T, method, path, apiKey string, body map[string]any) *http.Request { + t.Helper() + b, err := json.Marshal(body) + if err != nil { + t.Fatalf("marshal request body: %v", err) + } + req := httptest.NewRequest(method, path, bytes.NewReader(b)) + if apiKey != "" { + req.Header.Set("Authorization", "Bearer "+apiKey) + } + return req +} + +// recordingWriter is a minimal http.ResponseWriter/http.Flusher that logs +// the sequence of Write and Flush calls it receives, so a streaming test +// can assert flushes happen incrementally rather than once at the end. +type recordingWriter struct { + hdr http.Header + events []string + buf bytes.Buffer + status int +} + +func newRecordingWriter() *recordingWriter { + return &recordingWriter{hdr: http.Header{}} +} + +func (r *recordingWriter) Header() http.Header { return r.hdr } + +func (r *recordingWriter) WriteHeader(status int) { r.status = status } + +func (r *recordingWriter) Write(b []byte) (int, error) { + if r.status == 0 { + r.status = http.StatusOK // mirrors net/http's implicit-200 default when Write precedes any explicit WriteHeader + } + n, _ := r.buf.Write(b) + r.events = append(r.events, "write") + return n, nil +} + +func (r *recordingWriter) Flush() { r.events = append(r.events, "flush") } + +// alwaysErrStore is a counterStore stub whose every operation fails, used +// to drive the limiter's fail-closed (storeDown) path deterministically. +type alwaysErrStore struct{} + +var errStoreDownStub = errors.New("stub: store down") + +func (alwaysErrStore) incrBy(string, int64, time.Duration) (int64, error) { + return 0, errStoreDownStub +} + +func (alwaysErrStore) get(string) (int64, error) { + return 0, errStoreDownStub +} + +// TestHandleChat_HappyPath_NonStreaming_AccountsUsage drives a full +// chat-completion request through Gateway.ServeHTTP against a fake +// openai-shaped upstream, and asserts the response is forwarded verbatim +// and the reported usage lands in the limiter's own counters. +func TestHandleChat_HappyPath_NonStreaming_AccountsUsage(t *testing.T) { + const respBody = `{"id":"chatcmpl-1","object":"chat.completion","model":"gpt-real","choices":[{"index":0,"message":{"role":"assistant","content":"hi"},"finish_reason":"stop"}],"usage":{"prompt_tokens":10,"completion_tokens":5}}` + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(respBody)) + })) + defer srv.Close() + + cfg := CreateConfig() + cfg.Providers = map[string]*ProviderConfig{ + "openai": {Type: "openai", BaseURL: srv.URL, APIKey: "sk-up", Models: []string{"gpt-test"}}, + } + cfg.Groups = map[string]*GroupConfig{"default": {}} + cfg.Users = &UsersConfig{Inline: []*UserConfig{ + {Name: "alice", Group: "default", APIKey: "sk-alice", Limits: &LimitsConfig{}}, + }} + next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}) + h, err := New(context.Background(), next, cfg, "llmgw") + if err != nil { + t.Fatalf("New: %v", err) + } + gw, ok := h.(*Gateway) + if !ok { + t.Fatal("handler is not *Gateway") + } + + body := map[string]any{"model": "gpt-test", "messages": []any{map[string]any{"role": "user", "content": "hi"}}} + req := newUnifiedRequest(t, http.MethodPost, "/v1/chat/completions", "sk-alice", body) + rec := httptest.NewRecorder() + h.ServeHTTP(rec, req) + + if rec.Code != http.StatusOK { + t.Fatalf("status = %d, want 200, body=%s", rec.Code, rec.Body.String()) + } + if rec.Body.String() != respBody { + t.Errorf("body = %q, want verbatim upstream body %q", rec.Body.String(), respBody) + } + + tok, ok := gw.limiter.getCounter("user", "alice", metricTok, windowDay, time.Now()) + if !ok || tok != 15 { + t.Errorf("user token/day counter = %d (ok=%v), want 15", tok, ok) + } + reqCount, ok := gw.limiter.getCounter("user", "alice", metricReq, windowMin, time.Now()) + if !ok || reqCount != 1 { + t.Errorf("user request/min counter = %d (ok=%v), want 1", reqCount, ok) + } +} + +// TestHandleChat_Streaming_FlushesIncrementally proves a streaming chat +// completion is forwarded to the client chunk by chunk, flushing after +// each one, rather than buffered until the response completes. +func TestHandleChat_Streaming_FlushesIncrementally(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(http.StatusOK) + fl, ok := w.(http.Flusher) + if !ok { + t.Fatal("fake upstream ResponseWriter does not support Flush") + } + for _, c := range []string{ + `{"id":"c1","object":"chat.completion.chunk","choices":[{"delta":{"content":"Hel"}}]}`, + `{"id":"c1","object":"chat.completion.chunk","choices":[{"delta":{"content":"lo"}}]}`, + } { + _, _ = w.Write([]byte("data: " + c + "\n\n")) + fl.Flush() + } + _, _ = w.Write([]byte("data: [DONE]\n\n")) + fl.Flush() + })) + defer srv.Close() + + cfg := CreateConfig() + cfg.Providers = map[string]*ProviderConfig{ + "openai": {Type: "openai", BaseURL: srv.URL, APIKey: "sk-up", Models: []string{"gpt-test"}}, + } + cfg.Groups = map[string]*GroupConfig{"default": {}} + cfg.Users = &UsersConfig{Inline: []*UserConfig{{Name: "alice", Group: "default", APIKey: "sk-alice"}}} + next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}) + h, err := New(context.Background(), next, cfg, "llmgw") + if err != nil { + t.Fatalf("New: %v", err) + } + + body := map[string]any{"model": "gpt-test", "stream": true, "messages": []any{map[string]any{"role": "user", "content": "hi"}}} + req := newUnifiedRequest(t, http.MethodPost, "/v1/chat/completions", "sk-alice", body) + rw := newRecordingWriter() + h.ServeHTTP(rw, req) + + if rw.status != http.StatusOK { + t.Fatalf("status = %d, want 200", rw.status) + } + writes, flushes := 0, 0 + for _, e := range rw.events { + switch e { + case "write": + writes++ + case "flush": + flushes++ + } + } + if writes < 3 { + t.Fatalf("want >=3 writes (2 content chunks + [DONE]), got %d (events=%v)", writes, rw.events) + } + if flushes < writes { + t.Fatalf("want a flush per write (incremental delivery), got %d writes and %d flushes", writes, flushes) + } + if !strings.Contains(rw.buf.String(), "data: [DONE]") { + t.Errorf("body missing terminal [DONE], got %q", rw.buf.String()) + } +} + +// TestHandleChat_NoAPIKey_Returns401 asserts the unified routes require +// authentication before any resolve/limit work happens. +func TestHandleChat_NoAPIKey_Returns401(t *testing.T) { + cfg := CreateConfig() + cfg.Providers = map[string]*ProviderConfig{"openai": {Type: "openai", APIKey: "k"}} + next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}) + h, err := New(context.Background(), next, cfg, "llmgw") + if err != nil { + t.Fatalf("New: %v", err) + } + + req := newUnifiedRequest(t, http.MethodPost, "/v1/chat/completions", "", map[string]any{"model": "gpt-test"}) + rec := httptest.NewRecorder() + h.ServeHTTP(rec, req) + + if rec.Code != http.StatusUnauthorized { + t.Fatalf("status = %d, want 401, body=%s", rec.Code, rec.Body.String()) + } +} + +// TestHandleChat_DeniedModel_Returns403 covers a model that a configured +// provider knows, but the caller's group is not authorized to use. +func TestHandleChat_DeniedModel_Returns403(t *testing.T) { + cfg := CreateConfig() + cfg.Providers = map[string]*ProviderConfig{ + "openai": {Type: "openai", APIKey: "k", Models: []string{"gpt-test"}}, + } + cfg.Groups = map[string]*GroupConfig{"default": {Models: []string{"gpt-allowed-only"}}} + cfg.Users = &UsersConfig{Inline: []*UserConfig{{Name: "alice", Group: "default", APIKey: "sk-alice"}}} + next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}) + h, err := New(context.Background(), next, cfg, "llmgw") + if err != nil { + t.Fatalf("New: %v", err) + } + + body := map[string]any{"model": "gpt-test", "messages": []any{}} + req := newUnifiedRequest(t, http.MethodPost, "/v1/chat/completions", "sk-alice", body) + rec := httptest.NewRecorder() + h.ServeHTTP(rec, req) + + if rec.Code != http.StatusForbidden { + t.Fatalf("status = %d, want 403, body=%s", rec.Code, rec.Body.String()) + } +} + +// TestHandleChat_UnknownModel_Returns404 covers a model no configured +// provider knows at all. +func TestHandleChat_UnknownModel_Returns404(t *testing.T) { + cfg := CreateConfig() + cfg.Providers = map[string]*ProviderConfig{"openai": {Type: "openai", APIKey: "k", Models: []string{"gpt-test"}}} + cfg.Groups = map[string]*GroupConfig{"default": {}} + cfg.Users = &UsersConfig{Inline: []*UserConfig{{Name: "alice", Group: "default", APIKey: "sk-alice"}}} + next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}) + h, err := New(context.Background(), next, cfg, "llmgw") + if err != nil { + t.Fatalf("New: %v", err) + } + + body := map[string]any{"model": "no-such-model", "messages": []any{}} + req := newUnifiedRequest(t, http.MethodPost, "/v1/chat/completions", "sk-alice", body) + rec := httptest.NewRecorder() + h.ServeHTTP(rec, req) + + if rec.Code != http.StatusNotFound { + t.Fatalf("status = %d, want 404, body=%s", rec.Code, rec.Body.String()) + } +} + +// TestHandleChat_RequestLimitExceeded_Returns429WithRetryAfter drives two +// requests against a user limited to one per minute, and asserts the +// second is refused with a Retry-After header. +func TestHandleChat_RequestLimitExceeded_Returns429WithRetryAfter(t *testing.T) { + const respBody = `{"id":"c1","choices":[{"message":{"role":"assistant","content":"hi"}}],"usage":{"prompt_tokens":1,"completion_tokens":1}}` + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(respBody)) + })) + defer srv.Close() + + cfg := CreateConfig() + cfg.Providers = map[string]*ProviderConfig{ + "openai": {Type: "openai", BaseURL: srv.URL, APIKey: "k", Models: []string{"gpt-test"}}, + } + cfg.Groups = map[string]*GroupConfig{"default": {}} + cfg.Users = &UsersConfig{Inline: []*UserConfig{ + {Name: "alice", Group: "default", APIKey: "sk-alice", Limits: &LimitsConfig{RequestsPerMinute: 1}}, + }} + next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}) + h, err := New(context.Background(), next, cfg, "llmgw") + if err != nil { + t.Fatalf("New: %v", err) + } + + body := map[string]any{"model": "gpt-test", "messages": []any{}} + req1 := newUnifiedRequest(t, http.MethodPost, "/v1/chat/completions", "sk-alice", body) + rec1 := httptest.NewRecorder() + h.ServeHTTP(rec1, req1) + if rec1.Code != http.StatusOK { + t.Fatalf("first request status = %d, want 200, body=%s", rec1.Code, rec1.Body.String()) + } + + req2 := newUnifiedRequest(t, http.MethodPost, "/v1/chat/completions", "sk-alice", body) + rec2 := httptest.NewRecorder() + h.ServeHTTP(rec2, req2) + if rec2.Code != http.StatusTooManyRequests { + t.Fatalf("second request status = %d, want 429, body=%s", rec2.Code, rec2.Body.String()) + } + if ra := rec2.Header().Get("Retry-After"); ra == "" { + t.Error("want a Retry-After header on 429, got none") + } +} + +// TestHandleChat_StoreDown_FailClosed_Returns503 wires a counterStore +// stub whose every operation fails, with failOpen=false, and asserts the +// request is refused with 503 rather than silently enforcing limits +// against an unreachable backend. +func TestHandleChat_StoreDown_FailClosed_Returns503(t *testing.T) { + cfg := CreateConfig() + cfg.Providers = map[string]*ProviderConfig{"openai": {Type: "openai", APIKey: "k", Models: []string{"gpt-test"}}} + cfg.Groups = map[string]*GroupConfig{"default": {}} + cfg.Users = &UsersConfig{Inline: []*UserConfig{ + {Name: "alice", Group: "default", APIKey: "sk-alice", Limits: &LimitsConfig{RequestsPerMinute: 100}}, + }} + next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}) + h, err := New(context.Background(), next, cfg, "llmgw") + if err != nil { + t.Fatalf("New: %v", err) + } + gw, ok := h.(*Gateway) + if !ok { + t.Fatal("handler is not *Gateway") + } + gw.limiter = newLimiter(alwaysErrStore{}, false) + + body := map[string]any{"model": "gpt-test", "messages": []any{}} + req := newUnifiedRequest(t, http.MethodPost, "/v1/chat/completions", "sk-alice", body) + rec := httptest.NewRecorder() + h.ServeHTTP(rec, req) + + if rec.Code != http.StatusServiceUnavailable { + t.Fatalf("status = %d, want 503, body=%s", rec.Code, rec.Body.String()) + } +} + +// TestHandleEmbeddings_HappyPath mirrors the chat-completion happy path +// against POST /v1/embeddings. +func TestHandleEmbeddings_HappyPath(t *testing.T) { + const respBody = `{"object":"list","data":[{"embedding":[0.1,0.2]}],"model":"text-embedding-3","usage":{"prompt_tokens":5,"total_tokens":5}}` + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/v1/embeddings" { + t.Errorf("path = %q, want /v1/embeddings", r.URL.Path) + } + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(respBody)) + })) + defer srv.Close() + + cfg := CreateConfig() + cfg.Providers = map[string]*ProviderConfig{ + "openai": {Type: "openai", BaseURL: srv.URL, APIKey: "k", Models: []string{"text-embedding-3"}}, + } + cfg.Groups = map[string]*GroupConfig{"default": {}} + cfg.Users = &UsersConfig{Inline: []*UserConfig{{Name: "alice", Group: "default", APIKey: "sk-alice"}}} + next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}) + h, err := New(context.Background(), next, cfg, "llmgw") + if err != nil { + t.Fatalf("New: %v", err) + } + + body := map[string]any{"model": "text-embedding-3", "input": "hello"} + req := newUnifiedRequest(t, http.MethodPost, "/v1/embeddings", "sk-alice", body) + rec := httptest.NewRecorder() + h.ServeHTTP(rec, req) + + if rec.Code != http.StatusOK { + t.Fatalf("status = %d, want 200, body=%s", rec.Code, rec.Body.String()) + } + if rec.Body.String() != respBody { + t.Errorf("body = %q, want verbatim upstream body %q", rec.Body.String(), respBody) + } +} + +// TestHandleChat_AnthropicProviderPrefixedModel_EchoesAliasInResponse +// requests a provider-prefixed model id ("anthropic/claude-x"): the +// upstream call must use the bare upstream id, but the translated +// response's "model" field must echo back the exact id the client +// requested (ruling a, ALIAS ECHO), not the bare upstream form. +func TestHandleChat_AnthropicProviderPrefixedModel_EchoesAliasInResponse(t *testing.T) { + const anthResp = `{"id":"msg_01ABC","type":"message","role":"assistant","content":[{"type":"text","text":"hi there"}],"model":"claude-x","stop_reason":"end_turn","usage":{"input_tokens":10,"output_tokens":4}}` + + var gotBody []byte + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotBody, _ = io.ReadAll(r.Body) + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(anthResp)) + })) + defer srv.Close() + + cfg := CreateConfig() + cfg.Providers = map[string]*ProviderConfig{ + "anthropic": {Type: "anthropic", BaseURL: srv.URL, APIKey: "sk-ant", Models: []string{"claude-x"}}, + } + cfg.Groups = map[string]*GroupConfig{"default": {}} + cfg.Users = &UsersConfig{Inline: []*UserConfig{{Name: "alice", Group: "default", APIKey: "sk-alice"}}} + next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}) + h, err := New(context.Background(), next, cfg, "llmgw") + if err != nil { + t.Fatalf("New: %v", err) + } + + const requestedID = "anthropic/claude-x" + body := map[string]any{"model": requestedID, "messages": []any{map[string]any{"role": "user", "content": "hi"}}} + req := newUnifiedRequest(t, http.MethodPost, "/v1/chat/completions", "sk-alice", body) + rec := httptest.NewRecorder() + h.ServeHTTP(rec, req) + + if rec.Code != http.StatusOK { + t.Fatalf("status = %d, want 200, body=%s", rec.Code, rec.Body.String()) + } + + var out map[string]any + if err := json.Unmarshal(rec.Body.Bytes(), &out); err != nil { + t.Fatalf("decode response: %v", err) + } + if out["model"] != requestedID { + t.Errorf("response model = %v, want the client-requested alias %q echoed back", out["model"], requestedID) + } + + sentUpstream := decodeJSONBody(t, gotBody) + if sentUpstream["model"] != "claude-x" { + t.Errorf("upstream request model = %v, want bare upstream id %q", sentUpstream["model"], "claude-x") + } + if _, present := sentUpstream[gatewayAliasKey]; present { + t.Errorf("gatewayAliasKey %q leaked into the upstream request body: %v", gatewayAliasKey, sentUpstream) + } +} + +// TestHandleChat_UpstreamNon2xx_WrapsProviderErrorEnvelope asserts a +// non-2xx upstream response is passed through with its status code, but +// wrapped in the gateway's own envelope with the raw upstream body +// embedded under error.upstream (ruling c). +func TestHandleChat_UpstreamNon2xx_WrapsProviderErrorEnvelope(t *testing.T) { + const upstreamErrBody = `{"error":{"message":"rate limited upstream","type":"rate_limit_error"}}` + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusTooManyRequests) + _, _ = w.Write([]byte(upstreamErrBody)) + })) + defer srv.Close() + + cfg := CreateConfig() + cfg.Providers = map[string]*ProviderConfig{ + "openai": {Type: "openai", BaseURL: srv.URL, APIKey: "k", Models: []string{"gpt-test"}}, + } + cfg.Groups = map[string]*GroupConfig{"default": {}} + cfg.Users = &UsersConfig{Inline: []*UserConfig{{Name: "alice", Group: "default", APIKey: "sk-alice"}}} + next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}) + h, err := New(context.Background(), next, cfg, "llmgw") + if err != nil { + t.Fatalf("New: %v", err) + } + + body := map[string]any{"model": "gpt-test", "messages": []any{}} + req := newUnifiedRequest(t, http.MethodPost, "/v1/chat/completions", "sk-alice", body) + rec := httptest.NewRecorder() + h.ServeHTTP(rec, req) + + if rec.Code != http.StatusTooManyRequests { + t.Fatalf("status = %d, want 429, body=%s", rec.Code, rec.Body.String()) + } + + var out struct { + Error struct { + Upstream map[string]any `json:"upstream"` + Message string `json:"message"` + Type string `json:"type"` + Code string `json:"code"` + } `json:"error"` + } + if err := json.Unmarshal(rec.Body.Bytes(), &out); err != nil { + t.Fatalf("decode response: %v", err) + } + if out.Error.Type != "upstream_error" { + t.Errorf("error.type = %q, want upstream_error", out.Error.Type) + } + if out.Error.Code != "429" { + t.Errorf("error.code = %q, want %q", out.Error.Code, "429") + } + if out.Error.Upstream == nil { + t.Fatal("error.upstream is nil, want the parsed upstream body") + } + errObj, _ := out.Error.Upstream["error"].(map[string]any) + if errObj["message"] != "rate limited upstream" { + t.Errorf("error.upstream.error.message = %v, want %q", errObj["message"], "rate limited upstream") + } +} + +// TestHandleChat_MidStreamUpstreamDrop_NoTrailingEnvelope proves an +// adapter error arriving after the response has already started streaming +// gets logged, not turned into a second, conflicting envelope appended to +// a body the client already started receiving (ruling c). +func TestHandleChat_MidStreamUpstreamDrop_NoTrailingEnvelope(t *testing.T) { + const firstChunk = "data: {\"id\":\"c1\",\"object\":\"chat.completion.chunk\",\"choices\":[{\"delta\":{\"content\":\"Hi\"}}]}\n\n" + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(http.StatusOK) + fl, ok := w.(http.Flusher) + if !ok { + t.Fatal("fake upstream ResponseWriter does not support Flush") + } + _, _ = w.Write([]byte(firstChunk)) + fl.Flush() + + hj, ok := w.(http.Hijacker) + if !ok { + t.Fatal("fake upstream ResponseWriter does not support Hijack") + } + conn, _, hjErr := hj.Hijack() + if hjErr == nil { + _ = conn.Close() + } + })) + defer srv.Close() + + cfg := CreateConfig() + cfg.Providers = map[string]*ProviderConfig{ + "openai": {Type: "openai", BaseURL: srv.URL, APIKey: "k", Models: []string{"gpt-test"}}, + } + cfg.Groups = map[string]*GroupConfig{"default": {}} + cfg.Users = &UsersConfig{Inline: []*UserConfig{{Name: "alice", Group: "default", APIKey: "sk-alice"}}} + next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}) + h, err := New(context.Background(), next, cfg, "llmgw") + if err != nil { + t.Fatalf("New: %v", err) + } + + origStderr := os.Stderr + pr, pw, _ := os.Pipe() + os.Stderr = pw + defer func() { os.Stderr = origStderr }() + + body := map[string]any{"model": "gpt-test", "stream": true, "messages": []any{}} + req := newUnifiedRequest(t, http.MethodPost, "/v1/chat/completions", "sk-alice", body) + rec := httptest.NewRecorder() + h.ServeHTTP(rec, req) + + _ = pw.Close() // closing the pipe write end to unblock the read; error not actionable in a test + os.Stderr = origStderr + var logBuf bytes.Buffer + if _, err := io.Copy(&logBuf, pr); err != nil { + t.Fatalf("io.Copy: %v", err) + } + + if rec.Code != http.StatusOK { + t.Fatalf("status = %d, want 200 (headers already committed before the drop)", rec.Code) + } + if got := rec.Body.String(); got != firstChunk { + t.Errorf("body = %q, want exactly the one chunk written before the drop (no trailing error envelope)", got) + } + if !strings.Contains(logBuf.String(), "llmgw[llmgw] ERROR") { + t.Errorf("want the mid-stream error logged, got %q", logBuf.String()) + } +}