diff --git a/README.md b/README.md index d6eb701..6e5eda7 100644 --- a/README.md +++ b/README.md @@ -138,3 +138,15 @@ PII masking covers conversation text, tool results, assistant reasoning, tool-ca argument string values, and `user`. Tool definitions, names, IDs, schemas, metadata, provider options, and image/audio content are not scanned. Keep personal data out of those fields even when masking is enabled. + +Guest chat accounts (flagged `is_guest` by Metering; a landing-page visitor +with a small credit and no login) are ordinary accounts with two differences. +They can call only chat completions on models an operator ticked "Guests can +use this model" (`guest_allowed` from Metering); any other model or router is +refused with `guest_model_not_allowed` before routing, PII masking or Metering +run. And their chat requests always carry an output limit of at most +`GUEST_MAX_OUTPUT_TOKENS` (default 8192), written into the request so the +provider enforces it. Metering holds credit for that limit, so a guest's hold +fits its credit; other requests without `max_tokens` hold the model's maximum +output. An invalid value stops the gateway at startup. The Nexus repository's +`docs/guest-chat.md` describes guest chat. diff --git a/apps/gateway/guest_env_test.go b/apps/gateway/guest_env_test.go new file mode 100644 index 0000000..f20dd51 --- /dev/null +++ b/apps/gateway/guest_env_test.go @@ -0,0 +1,20 @@ +package main + +import "testing" + +func TestPositiveIntEnv(t *testing.T) { + t.Setenv("GUEST_MAX_OUTPUT_TOKENS", "") + if n, err := positiveIntEnv("GUEST_MAX_OUTPUT_TOKENS", 8192); err != nil || n != 8192 { + t.Fatalf("unset = %d, %v; want the default", n, err) + } + t.Setenv("GUEST_MAX_OUTPUT_TOKENS", "4096") + if n, err := positiveIntEnv("GUEST_MAX_OUTPUT_TOKENS", 8192); err != nil || n != 4096 { + t.Fatalf("4096 = %d, %v", n, err) + } + for _, bad := range []string{"0", "-1", "8k"} { + t.Setenv("GUEST_MAX_OUTPUT_TOKENS", bad) + if _, err := positiveIntEnv("GUEST_MAX_OUTPUT_TOKENS", 8192); err == nil { + t.Fatalf("%q was accepted", bad) + } + } +} diff --git a/apps/gateway/internal/adapters/http/handlers/chat_completions_handler.go b/apps/gateway/internal/adapters/http/handlers/chat_completions_handler.go index 2644971..65b9516 100644 --- a/apps/gateway/internal/adapters/http/handlers/chat_completions_handler.go +++ b/apps/gateway/internal/adapters/http/handlers/chat_completions_handler.go @@ -1,11 +1,9 @@ package handlers import ( - "context" "io" "net/http" "sync" - "sync/atomic" "time" "github.com/dappnode/dappnode-nexus-gateway/apps/gateway/internal/adapters/http/mapper" @@ -16,10 +14,6 @@ import ( "github.com/dappnode/dappnode-nexus-gateway/pkg/domain" ) -// Bound the wait for usage after generation has finished, without extending -// the lifetime of requests canceled before completion. -const streamUsageDrainTimeout = 5 * time.Second - // ChatCompletionsHandler handles POST /v1/chat/completions. type ChatCompletionsHandler struct { service *services.ChatCompletionsService @@ -50,27 +44,16 @@ func (h *ChatCompletionsHandler) Handle(w http.ResponseWriter, r *http.Request) } // The gateway is a proxy: the request goes upstream as the client sent it - // and the provider's response comes back as it came. - upstream, cancelUpstream := context.WithCancel(context.WithoutCancel(r.Context())) - defer cancelUpstream() - var finished atomic.Bool - stopWatch := context.AfterFunc(r.Context(), func() { - // A client gone after the finish reason still gets its usage metered: - // the upstream is drained briefly instead of cut off. - if finished.Load() { - time.AfterFunc(streamUsageDrainTimeout, cancelUpstream) - return - } - cancelUpstream() - }) - defer stopWatch() - call, err := h.service.Proxy(upstream, body, genReq, token) + // and the provider's response comes back as it came. Once the provider + // runs, a client that hangs up no longer stops it (see Proxy): the relay + // keeps reading to the end, where the usage to charge arrives. + call, err := h.service.Proxy(r.Context(), body, genReq, token) if err != nil { WriteErrorWithLog(w, r, h.logger, err) return } if call.Stream != nil { - h.relay(w, r, call.Stream, &finished) + h.relay(w, r, call.Stream) return } w.Header().Set("Content-Type", "application/json") @@ -81,7 +64,7 @@ func (h *ChatCompletionsHandler) Handle(w http.ResponseWriter, r *http.Request) // relay sends a provider's stream to the client line by line, as it came. // Keep-alive comments fill long silences (slow reasoning), and a stream the // provider finished without [DONE] gets one, so every client sees an end. -func (h *ChatCompletionsHandler) relay(w http.ResponseWriter, r *http.Request, stream *services.ProxyStream, finished *atomic.Bool) { +func (h *ChatCompletionsHandler) relay(w http.ResponseWriter, r *http.Request, stream *services.ProxyStream) { defer stream.Close() sw, err := sse.NewWriter(w) if err != nil { @@ -124,9 +107,6 @@ func (h *ChatCompletionsHandler) relay(w http.ResponseWriter, r *http.Request, s readErr = err break } - if stream.Finished() { - finished.Store(true) - } write(line) } switch { diff --git a/apps/gateway/internal/adapters/http/handlers/chat_completions_stream_test.go b/apps/gateway/internal/adapters/http/handlers/chat_completions_stream_test.go index d0e14d1..f1a4a89 100644 --- a/apps/gateway/internal/adapters/http/handlers/chat_completions_stream_test.go +++ b/apps/gateway/internal/adapters/http/handlers/chat_completions_stream_test.go @@ -5,7 +5,6 @@ import ( "io" "net/http/httptest" "strings" - "sync/atomic" "testing" "github.com/dappnode/dappnode-nexus-gateway/apps/gateway/internal/application/services" @@ -45,8 +44,7 @@ func TestRelay_EndsEveryStream(t *testing.T) { t.Run(name, func(t *testing.T) { recorder := httptest.NewRecorder() handler := &ChatCompletionsHandler{logger: &confidentialTestLogger{}} - var finished atomic.Bool - handler.relay(recorder, httptest.NewRequest("POST", "/v1/chat/completions", nil), services.NewProxyStream(brokenBody{tc.body}, "m"), &finished) + handler.relay(recorder, httptest.NewRequest("POST", "/v1/chat/completions", nil), services.NewProxyStream(brokenBody{tc.body}, "m")) body := recorder.Body.String() for _, want := range tc.want { if !strings.Contains(body, want) { diff --git a/apps/gateway/internal/adapters/http/handlers/stream_usage_repro_test.go b/apps/gateway/internal/adapters/http/handlers/stream_usage_repro_test.go index dbf6190..1359cd2 100644 --- a/apps/gateway/internal/adapters/http/handlers/stream_usage_repro_test.go +++ b/apps/gateway/internal/adapters/http/handlers/stream_usage_repro_test.go @@ -49,13 +49,14 @@ func TestStreamingUsageRetention(t *testing.T) { for _, tc := range []struct { name string - placement string closeOnDone bool closeOnDelta bool wantUsage bool }{ - {name: "trailing_usage", placement: "trailing", wantUsage: true}, - {name: "client_cancels_before_finish", placement: "unfinished", closeOnDelta: true}, + {name: "trailing_usage", wantUsage: true}, + // A client that hangs up mid-answer doesn't stop the provider: the + // answer finishes and its usage is charged, as for any other request. + {name: "client_hangs_up_before_finish", closeOnDelta: true, wantUsage: true}, } { t.Run(tc.name, func(t *testing.T) { ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) @@ -79,13 +80,14 @@ func TestStreamingUsageRetention(t *testing.T) { w.(http.Flusher).Flush() } emit(toolDelta, "null", "null") - if tc.placement == "unfinished" { + if tc.closeOnDelta { + // Let the client hang up before the answer finishes. select { + case <-time.After(50 * time.Millisecond): case <-r.Context().Done(): close(upstreamCanceled) - case <-ctx.Done(): + return } - return } emit(`{}`, `"tool_calls"`, "null") // Delay trailing usage so the handler's drain loop has work to do. @@ -162,17 +164,6 @@ func TestStreamingUsageRetention(t *testing.T) { } select { case recorded := <-meter.recorded: - if tc.closeOnDelta { - if recorded.err == nil { - t.Fatal("client cancellation before finish was recorded as success") - } - select { - case <-upstreamCanceled: - case <-ctx.Done(): - t.Fatal("upstream continued generating after early client cancellation") - } - return - } if recorded.err != nil { t.Fatalf("recorded failure: %v", recorded.err) } diff --git a/apps/gateway/internal/adapters/metering/client.go b/apps/gateway/internal/adapters/metering/client.go index f95852e..6407457 100644 --- a/apps/gateway/internal/adapters/metering/client.go +++ b/apps/gateway/internal/adapters/metering/client.go @@ -58,6 +58,7 @@ func (c *Client) Reserve(ctx context.Context, auth domain.AuthContext, endpoint FallbackUsed: req.FallbackUsed, Stream: req.Stream, MaxOutputTokens: req.MaxOutputTokens, + EstimatedInputTokens: req.EstimatedInputTokens, } var res reserveResponse sent := time.Now() // A lease measured from here never outlasts Metering's. @@ -420,6 +421,7 @@ type reserveRequest struct { FallbackUsed *bool `json:"fallback_used,omitempty"` Stream bool `json:"stream"` MaxOutputTokens *int `json:"max_output_tokens,omitempty"` + EstimatedInputTokens *int64 `json:"estimated_input_tokens,omitempty"` } type reserveResponse struct { diff --git a/apps/gateway/internal/adapters/metering/client_test.go b/apps/gateway/internal/adapters/metering/client_test.go index c7c30fa..bd19297 100644 --- a/apps/gateway/internal/adapters/metering/client_test.go +++ b/apps/gateway/internal/adapters/metering/client_test.go @@ -38,6 +38,37 @@ func TestClientReserve_MapsInsufficientBalance(t *testing.T) { } } +func TestClientReserve_SendsHoldSizing(t *testing.T) { + var payload map[string]any + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if err := json.NewDecoder(r.Body).Decode(&payload); err != nil { + t.Fatalf("decode reserve payload: %v", err) + } + _, _ = w.Write([]byte(`{"reservation_id":"r-1"}`)) + })) + defer server.Close() + + estimate, maxTokens := int64(9_216), 512 + client := NewClient(server.URL, "token", time.Second) + req := domain.GenerateRequest{PublicModelID: "model", EstimatedInputTokens: &estimate, MaxOutputTokens: &maxTokens} + if _, err := client.Reserve(context.Background(), testAuth(), domain.EndpointChatCompletions, req, testModel(), "gw-1"); err != nil { + t.Fatalf("Reserve: %v", err) + } + if payload["estimated_input_tokens"] != float64(9_216) || payload["max_output_tokens"] != float64(512) { + t.Fatalf("hold sizing = estimate %v/max output %v, want 9216/512", payload["estimated_input_tokens"], payload["max_output_tokens"]) + } + + // Without an estimate the field is omitted, so Metering holds the whole + // context window. + payload = nil + if _, err := client.Reserve(context.Background(), testAuth(), domain.EndpointChatCompletions, domain.GenerateRequest{PublicModelID: "model"}, testModel(), "gw-2"); err != nil { + t.Fatalf("Reserve without estimate: %v", err) + } + if _, present := payload["estimated_input_tokens"]; present { + t.Fatalf("payload = %v, want no estimated_input_tokens", payload) + } +} + func TestClientReserve_MapsInternalContractErrorsToInternalError(t *testing.T) { server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { w.WriteHeader(http.StatusConflict) diff --git a/apps/gateway/internal/adapters/metering/runtime_client.go b/apps/gateway/internal/adapters/metering/runtime_client.go index ba39c2d..ed7f7d4 100644 --- a/apps/gateway/internal/adapters/metering/runtime_client.go +++ b/apps/gateway/internal/adapters/metering/runtime_client.go @@ -33,6 +33,7 @@ func (c *Client) AuthenticateAPIKey(ctx context.Context, rawKey string) (domain. ID: response.Account.ID, ExternalCustomerID: response.Account.ExternalCustomerID, Status: response.Account.Status, + IsGuest: response.Account.IsGuest, }, APIKey: domain.APIKey{ ID: response.APIKey.ID, @@ -357,6 +358,7 @@ func (m runtimePublicModel) toDomain() domain.PublicModel { OutputPricePerMillion: priceFromMicrocents(*m.OutputPricePer1MTokensMicrocents), Currency: m.Currency, Active: m.Active, + GuestAllowed: m.GuestAllowed, } if m.Fallback != nil { model.Fallback = &domain.ProviderTarget{ @@ -413,6 +415,7 @@ type runtimeAccount struct { ID string `json:"id"` ExternalCustomerID *string `json:"external_customer_id,omitempty"` Status string `json:"status"` + IsGuest bool `json:"is_guest"` } type runtimeAPIKey struct { @@ -451,6 +454,7 @@ type runtimePublicModel struct { CacheWritePricePer1MTokensMicrocents *int64 `json:"cache_write_price_per_1m_tokens_microcents,omitempty"` Currency string `json:"currency"` Active bool `json:"active"` + GuestAllowed bool `json:"guest_allowed"` } type runtimeProviderTarget struct { diff --git a/apps/gateway/internal/adapters/metering/runtime_client_test.go b/apps/gateway/internal/adapters/metering/runtime_client_test.go index 2ad1ca5..5f20ce8 100644 --- a/apps/gateway/internal/adapters/metering/runtime_client_test.go +++ b/apps/gateway/internal/adapters/metering/runtime_client_test.go @@ -52,6 +52,32 @@ func TestRuntimeClientAuthenticateHashesAPIKeyLocally(t *testing.T) { } } +// The guest flag decides whether the guest output limit applies, so it must +// survive the trip from metering; an older metering that omits it means "not +// a guest". +func TestRuntimeClientAuthenticateReadsTheGuestFlag(t *testing.T) { + for _, tc := range []struct { + account string + want bool + }{ + {`{"id":"acc-1","status":"active","is_guest":true}`, true}, + {`{"id":"acc-1","status":"active","is_guest":false}`, false}, + {`{"id":"acc-1","status":"active"}`, false}, + } { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + _, _ = w.Write([]byte(`{"account":` + tc.account + `,"api_key":{"id":"key-1","account_id":"acc-1","key_prefix":"nxs_1234","active":true,"pii_mode":"off"}}`)) + })) + auth, err := NewClient(server.URL, "token", time.Second).AuthenticateAPIKey(context.Background(), "sk-x") + server.Close() + if err != nil { + t.Fatalf("AuthenticateAPIKey: %v", err) + } + if auth.Account.IsGuest != tc.want { + t.Fatalf("account %s: IsGuest = %v, want %v", tc.account, auth.Account.IsGuest, tc.want) + } + } +} + func TestRuntimeClientPreservesUserAuthenticationErrors(t *testing.T) { tests := []struct { name string diff --git a/apps/gateway/internal/adapters/pii/presidio/adapter.go b/apps/gateway/internal/adapters/pii/presidio/adapter.go index fbf6e16..82720e5 100644 --- a/apps/gateway/internal/adapters/pii/presidio/adapter.go +++ b/apps/gateway/internal/adapters/pii/presidio/adapter.go @@ -231,17 +231,18 @@ func convertEntities(text string, items []analyzeResponseItem) []domain.PIIEntit return nil } + // totalRunes lets us bounds-check character offsets cheaply. + totalRunes := utf8.RuneCountInString(text) + // Build a runeIndex -> byteIndex lookup table. The final entry maps the // past-the-end rune index to len(text), so we can convert exclusive ends. - runes := make([]int, 0, len(text)+1) + // The capacity is a hint; append grows the slice for the final entry. + runes := make([]int, 0, totalRunes) for i := range text { // i is the byte index at each rune start runes = append(runes, i) } runes = append(runes, len(text)) - // totalRunes lets us bounds-check character offsets cheaply. - totalRunes := utf8.RuneCountInString(text) - out := make([]domain.PIIEntity, 0, len(items)) for _, it := range items { if it.Start < 0 || it.End < it.Start || it.End > totalRunes { diff --git a/apps/gateway/internal/application/services/generate_service.go b/apps/gateway/internal/application/services/generate_service.go index feeac0f..6db3396 100644 --- a/apps/gateway/internal/application/services/generate_service.go +++ b/apps/gateway/internal/application/services/generate_service.go @@ -26,6 +26,8 @@ type GenerateService struct { piiFailOpen bool tinfoilProofs ports.TinfoilTransportProofRepository logger ports.Logger + // guestMaxOutputTokens is the output limit of guest chat requests. + guestMaxOutputTokens int } func NewGenerateService( @@ -46,6 +48,8 @@ func NewGenerateService( pii: pii, piiLang: "en", logger: logger, + + guestMaxOutputTokens: DefaultGuestMaxOutputTokens, } } @@ -58,6 +62,12 @@ func (s *GenerateService) SetPIIOptions(language string, failOpen bool) { s.piiFailOpen = failOpen } +// SetGuestMaxOutputTokens sets the output limit of guest chat requests; see +// capOutput. +func (s *GenerateService) SetGuestMaxOutputTokens(limit int) { + s.guestMaxOutputTokens = limit +} + func (s *GenerateService) SetTinfoilProofRepository(proofs ports.TinfoilTransportProofRepository) { s.tinfoilProofs = proofs } diff --git a/apps/gateway/internal/application/services/guest_models.go b/apps/gateway/internal/application/services/guest_models.go new file mode 100644 index 0000000..5f8ca9d --- /dev/null +++ b/apps/gateway/internal/application/services/guest_models.go @@ -0,0 +1,30 @@ +package services + +import ( + "context" + + "github.com/dappnode/dappnode-nexus-gateway/pkg/domain" +) + +// checkGuestModel refuses a guest key any request but chat on a model an +// operator ticked "Guests can use this model" (public_models.guest_allowed). +// It runs before routing, PII masking and Metering: a router can call a paid +// API (Jev calls TypeSafe) before credit is checked, and a free model would +// cost a guest nothing against its credit. Routers are never public models, +// so they are always refused. +func (s *GenerateService) checkGuestModel(ctx context.Context, endpoint, publicModelID string) error { + if endpoint != domain.EndpointChatCompletions { + return domain.ErrGuestModelNotAllowed(publicModelID) + } + model, err := s.catalog.GetPublicModel(ctx, publicModelID) + if isUnsupportedModelError(err) { + return domain.ErrGuestModelNotAllowed(publicModelID) + } + if err != nil { + return err + } + if !model.GuestAllowed { + return domain.ErrGuestModelNotAllowed(publicModelID) + } + return nil +} diff --git a/apps/gateway/internal/application/services/guest_models_test.go b/apps/gateway/internal/application/services/guest_models_test.go new file mode 100644 index 0000000..14616dc --- /dev/null +++ b/apps/gateway/internal/application/services/guest_models_test.go @@ -0,0 +1,79 @@ +package services + +import ( + "context" + "errors" + "testing" + + "github.com/dappnode/dappnode-nexus-gateway/apps/gateway/internal/adapters/http/mapper" + "github.com/dappnode/dappnode-nexus-gateway/pkg/domain" +) + +// A guest key reaches only models ticked for guests. Anything else is refused +// before routing, PII masking, Metering or a provider runs: a router can call +// a paid API, and a free model would cost nothing against the guest's credit. +func TestProxy_GuestKeysCallOnlyGuestModels(t *testing.T) { + guestModel := domain.PublicModel{ + PublicModelID: "flash", ProviderModelID: "flash", ProviderConfig: domain.ProviderConfig{ProviderName: "openai"}, + SupportsChatCompletions: true, MaxContextWindow: 1000, MaxOutputTokens: 100, GuestAllowed: true, + } + freeModel := guestModel + freeModel.PublicModelID, freeModel.GuestAllowed = "free", false + + for name, tc := range map[string]struct { + endpoint, model string + wantErr string + }{ + "ticked model": {domain.EndpointChatCompletions, "flash", ""}, + "model not ticked": {domain.EndpointChatCompletions, "free", domain.ErrCodeGuestModelNotAllowed}, + "router": {domain.EndpointChatCompletions, "nexus/auto", domain.ErrCodeGuestModelNotAllowed}, + "other endpoint": {domain.EndpointSystemOne, "flash", domain.ErrCodeGuestModelNotAllowed}, + "unknown model or route": {domain.EndpointChatCompletions, "nope", domain.ErrCodeGuestModelNotAllowed}, + } { + t.Run(name, func(t *testing.T) { + meter, provider, router := &stubUsageMeter{}, &stubProvider{}, &stubRouterClient{ + decision: domain.RouteDecision{PublicModelID: "flash"}, + } + svc := newGuestTestService(meter, provider, true) + svc.catalog = &stubModelCatalog{ + models: map[string]domain.PublicModel{"flash": guestModel, "free": freeModel}, + routers: map[string]domain.RouterEntry{"nexus/auto": {RouterID: "nexus/auto"}}, + } + svc.router = router + body := `{"model":"` + tc.model + `","messages":[{"role":"user","content":"hi"}]}` + req, err := mapper.ChatCompletionRequestToDomain([]byte(body)) + if err != nil { + t.Fatalf("map request: %v", err) + } + _, err = svc.Proxy(context.Background(), tc.endpoint, []byte(body), req, "sk-test") + if tc.wantErr == "" { + if err != nil { + t.Fatalf("Proxy: %v", err) + } + return + } + var gwErr *domain.GatewayError + if !errors.As(err, &gwErr) || gwErr.Code != tc.wantErr || gwErr.HTTPStatus != 403 { + t.Fatalf("err = %v, want 403 %s", err, tc.wantErr) + } + if router.calls != 0 || meter.reserveCalls != 0 || provider.calls != 0 { + t.Fatalf("router %d, reserve %d, provider %d calls; want none", router.calls, meter.reserveCalls, provider.calls) + } + }) + } +} + +// When the catalog can't answer, a guest gets the catalog's error, not a +// misleading "not allowed". +func TestProxy_GuestModelCheckPassesCatalogErrors(t *testing.T) { + meter, provider := &stubUsageMeter{}, &stubProvider{} + svc := newGuestTestService(meter, provider, true) + svc.catalog = &stubModelCatalog{err: domain.ErrInternal("metering unavailable")} + body := `{"model":"flash","messages":[{"role":"user","content":"hi"}]}` + req, _ := mapper.ChatCompletionRequestToDomain([]byte(body)) + _, err := svc.Proxy(context.Background(), domain.EndpointChatCompletions, []byte(body), req, "sk-test") + var gwErr *domain.GatewayError + if !errors.As(err, &gwErr) || gwErr.Code != domain.ErrCodeInternalError { + t.Fatalf("err = %v, want internal_error", err) + } +} diff --git a/apps/gateway/internal/application/services/input_estimate.go b/apps/gateway/internal/application/services/input_estimate.go new file mode 100644 index 0000000..a4e7caf --- /dev/null +++ b/apps/gateway/internal/application/services/input_estimate.go @@ -0,0 +1,79 @@ +package services + +import ( + "bytes" + "encoding/json" + + "github.com/dappnode/dappnode-nexus-gateway/pkg/domain" +) + +// Metering holds credit for a request's prompt before the provider runs it. +// The hold needs an upper bound, not a guess: the prompt is charged at its +// real size when the request settles, and an estimate that is too low lets a +// request start that the account may not be able to pay for. +// +// The bound is the body's size in bytes. The tokenizers behind the models +// Nexus serves (byte-level BPE, and SentencePiece with byte fallback) never +// produce more tokens than the UTF-8 bytes of the text they encode. The JSON +// body carries every prompt string at least once, plus quotes, keys and +// escapes, so its length also covers the role markers the chat template adds +// per message. promptOverheadTokens covers what a provider adds that is not in +// the body: rendered tool instructions, a default system prompt, the upstream +// model name. +// +// Media is not text. An image part counts imageTokens instead of its bytes: an +// inline base64 image is far longer in bytes than in tokens, and a linked one +// far shorter. Audio, video, files, and part types the gateway does not know +// have no bound it can compute, so those requests get no estimate and the hold +// covers the model's whole context window, as it did before. +const ( + // imageTokens is above what any model Nexus serves spends on one image. + imageTokens = 16_384 + promptOverheadTokens = 1_024 +) + +// estimateInputTokens returns an upper bound on the prompt tokens of body, +// the request exactly as it goes upstream. ok is false when there is no bound. +func estimateInputTokens(endpoint string, body []byte) (tokens int64, ok bool) { + // SystemOne's provider builds its own prompt from the request, so its + // size is unknown here. + if endpoint != domain.EndpointChatCompletions { + return 0, false + } + var req struct { + Messages []struct { + Content json.RawMessage `json:"content"` + } `json:"messages"` + } + if err := json.Unmarshal(body, &req); err != nil { + return 0, false + } + tokens = int64(len(body)) + promptOverheadTokens + for _, message := range req.Messages { + content := bytes.TrimSpace(message.Content) + if len(content) == 0 || content[0] != '[' { + continue // A string or null: text, already counted. + } + var parts []json.RawMessage + if err := json.Unmarshal(content, &parts); err != nil { + return 0, false + } + for _, part := range parts { + var p struct { + Type string `json:"type"` + } + if err := json.Unmarshal(part, &p); err != nil { + return 0, false + } + switch p.Type { + case "text", "input_text", "output_text", "refusal": + // Text, already counted in the body's bytes. + case "image_url", "image", "input_image": + tokens += imageTokens - int64(len(part)) + default: + return 0, false + } + } + } + return max(tokens, 1), true +} diff --git a/apps/gateway/internal/application/services/input_estimate_test.go b/apps/gateway/internal/application/services/input_estimate_test.go new file mode 100644 index 0000000..f5298da --- /dev/null +++ b/apps/gateway/internal/application/services/input_estimate_test.go @@ -0,0 +1,146 @@ +package services + +import ( + "context" + "strings" + "testing" + + "github.com/dappnode/dappnode-nexus-gateway/pkg/domain" +) + +func TestEstimateInputTokens(t *testing.T) { + inlineImage := `{"type":"image_url","image_url":{"url":"data:image/png;base64,` + strings.Repeat("A", 200_000) + `"}}` + linkedImage := `{"type":"image_url","image_url":{"url":"https://example.com/cat.png"}}` + text := `{"type":"text","text":"what is this?"}` + + cases := []struct { + name string + body string + // extra is what the estimate adds to the body's bytes; ok=false means + // no bound. + extra func(body string) int64 + ok bool + }{ + { + name: "string content counts every byte", + body: `{"model":"m","messages":[{"role":"system","content":"be brief"},{"role":"user","content":"hello"}]}`, + extra: func(string) int64 { return 0 }, + ok: true, + }, + { + name: "text parts count every byte", + body: `{"model":"m","messages":[{"role":"user","content":[` + text + `,{"type":"input_text","text":"x"}]}]}`, + extra: func(string) int64 { return 0 }, + ok: true, + }, + { + name: "null content and tool calls", + body: `{"model":"m","messages":[{"role":"assistant","content":null,"tool_calls":[{"id":"1","type":"function","function":{"name":"f","arguments":"{}"}}]},{"role":"tool","tool_call_id":"1","content":"ok"}]}`, + extra: func(string) int64 { return 0 }, + ok: true, + }, + { + name: "inline image counts a fixed allowance instead of its bytes", + body: `{"model":"m","messages":[{"role":"user","content":[` + text + `,` + inlineImage + `]}]}`, + extra: func(string) int64 { return imageTokens - int64(len(inlineImage)) }, + ok: true, + }, + { + name: "linked image counts the allowance too", + body: `{"model":"m","messages":[{"role":"user","content":[` + linkedImage + `]}]}`, + extra: func(string) int64 { return imageTokens - int64(len(linkedImage)) }, + ok: true, + }, + { + name: "every image counts", + body: `{"model":"m","messages":[{"role":"user","content":[` + linkedImage + `]},{"role":"tool","tool_call_id":"1","content":[` + linkedImage + `]}]}`, + extra: func(string) int64 { return 2 * (imageTokens - int64(len(linkedImage))) }, + ok: true, + }, + { + name: "audio has no bound", + body: `{"model":"m","messages":[{"role":"user","content":[{"type":"input_audio","input_audio":{"data":"AAAA","format":"wav"}}]}]}`, + }, + { + name: "video has no bound", + body: `{"model":"m","messages":[{"role":"user","content":[{"type":"video_url","video_url":{"url":"https://example.com/a.mp4"}}]}]}`, + }, + { + name: "files have no bound", + body: `{"model":"m","messages":[{"role":"user","content":[{"type":"file","file":{"file_id":"f-1"}}]}]}`, + }, + { + name: "an unknown part type has no bound", + body: `{"model":"m","messages":[{"role":"user","content":[{"type":"hologram"}]}]}`, + }, + { + name: "a part without a type has no bound", + body: `{"model":"m","messages":[{"role":"user","content":[{"text":"hi"}]}]}`, + }, + { + name: "a part that is not an object has no bound", + body: `{"model":"m","messages":[{"role":"user","content":["hi"]}]}`, + }, + { + name: "invalid JSON has no bound", + body: `{"model":"m","messages":[`, + }, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + got, ok := estimateInputTokens(domain.EndpointChatCompletions, []byte(tc.body)) + if ok != tc.ok { + t.Fatalf("ok = %v, want %v", ok, tc.ok) + } + if !ok { + return + } + want := int64(len(tc.body)) + promptOverheadTokens + tc.extra(tc.body) + if got != want { + t.Fatalf("estimate = %d, want %d", got, want) + } + }) + } +} + +func TestEstimateInputTokens_InlineImageDoesNotHoldItsBytes(t *testing.T) { + image := `{"type":"image_url","image_url":{"url":"data:image/jpeg;base64,` + strings.Repeat("A", 2_000_000) + `"}}` + body := `{"model":"m","messages":[{"role":"user","content":[` + image + `]}]}` + got, ok := estimateInputTokens(domain.EndpointChatCompletions, []byte(body)) + if !ok || got > imageTokens+promptOverheadTokens+100 { + t.Fatalf("estimate = %d (ok=%v), want about one image allowance, not 2M bytes", got, ok) + } +} + +func TestEstimateInputTokens_SystemOneHasNoBound(t *testing.T) { + if _, ok := estimateInputTokens(domain.EndpointSystemOne, []byte(`{"model":"m","input":"hi"}`)); ok { + t.Fatal("systemone builds its own prompt, so it must not get an estimate") + } +} + +func TestProxy_SendsThePromptEstimateWithTheHold(t *testing.T) { + meter := &stubUsageMeter{} + svc := newDirectModelGenerateService(meter, &stubProvider{}) + body := `{"model":"model","messages":[{"role":"user","content":"hello"}]}` + + if _, err := svc.Proxy(context.Background(), domain.EndpointChatCompletions, []byte(body), domain.GenerateRequest{PublicModelID: "model"}, "sk-test"); err != nil { + t.Fatalf("Proxy: %v", err) + } + estimate := meter.lastReserveReq.EstimatedInputTokens + if want := int64(len(body)) + promptOverheadTokens; estimate == nil || *estimate != want { + t.Fatalf("estimate sent with the hold = %v, want %d", estimate, want) + } +} + +func TestProxy_SendsNoEstimateWhenThePromptHasNoBound(t *testing.T) { + meter := &stubUsageMeter{} + svc := newDirectModelGenerateService(meter, &stubProvider{}) + body := `{"model":"model","messages":[{"role":"user","content":[{"type":"input_audio","input_audio":{"data":"AAAA","format":"wav"}}]}]}` + + if _, err := svc.Proxy(context.Background(), domain.EndpointChatCompletions, []byte(body), domain.GenerateRequest{PublicModelID: "model"}, "sk-test"); err != nil { + t.Fatalf("Proxy: %v", err) + } + if meter.lastReserveReq.EstimatedInputTokens != nil { + t.Fatalf("estimate = %d, want none so the hold covers the context window", *meter.lastReserveReq.EstimatedInputTokens) + } +} diff --git a/apps/gateway/internal/application/services/output_limit.go b/apps/gateway/internal/application/services/output_limit.go new file mode 100644 index 0000000..137f4b7 --- /dev/null +++ b/apps/gateway/internal/application/services/output_limit.go @@ -0,0 +1,47 @@ +package services + +import ( + "bytes" + "encoding/json" + + "github.com/dappnode/dappnode-nexus-gateway/pkg/domain" +) + +// DefaultGuestMaxOutputTokens is the output limit of a guest's chat request +// unless GUEST_MAX_OUTPUT_TOKENS sets another. Without it a guest's request +// would hold the model's maximum output, far more than a guest's credit. It +// still covers almost every answer, reasoning included: over a week of production traffic on DeepSeek V4.1 Flash, answers +// averaged about 1,100 output tokens and the longest was about 15,000 +// (2026-11-04). +const DefaultGuestMaxOutputTokens = 8192 + +// capOutput writes an output limit into a guest's chat request: the guest's +// own limit when it is smaller, otherwise limit. Metering holds credit for +// that limit and the provider enforces it. Other accounts' requests are left +// as sent: without max_tokens, Metering holds the model's maximum output. +func capOutput(raw []byte, req domain.GenerateRequest, limit int) ([]byte, domain.GenerateRequest, error) { + if req.MaxOutputTokens != nil && *req.MaxOutputTokens > 0 && *req.MaxOutputTokens <= limit { + return raw, req, nil + } + dec := json.NewDecoder(bytes.NewReader(raw)) + dec.UseNumber() // Forward numbers exactly as sent. + var body map[string]any + if err := dec.Decode(&body); err != nil || body == nil { + return nil, req, domain.ErrInvalidField("invalid JSON body") + } + // Keep the client's field name, or use OpenAI's current one. The provider + // adapter sends max_tokens to providers that ignore max_completion_tokens + // (DeepSeek); every other provider enforces both (checked 2026-11-04). + // The mapper refuses requests that set both. + if _, ok := body["max_tokens"]; ok { + body["max_tokens"] = limit + } else { + body["max_completion_tokens"] = limit + } + rewritten, err := json.Marshal(body) + if err != nil { + return nil, req, domain.ErrInternal("an internal error occurred").WithMeta("reason", "re-encode output limit", "error", err.Error()) + } + req.MaxOutputTokens = &limit + return rewritten, req, nil +} diff --git a/apps/gateway/internal/application/services/output_limit_test.go b/apps/gateway/internal/application/services/output_limit_test.go new file mode 100644 index 0000000..5ef3fba --- /dev/null +++ b/apps/gateway/internal/application/services/output_limit_test.go @@ -0,0 +1,96 @@ +package services + +import ( + "context" + "encoding/json" + "testing" + + "github.com/dappnode/dappnode-nexus-gateway/apps/gateway/internal/adapters/http/mapper" + "github.com/dappnode/dappnode-nexus-gateway/pkg/domain" +) + +func newGuestTestService(meter *stubUsageMeter, provider *stubProvider, isGuest bool) *GenerateService { + svc := newDirectModelGenerateService(meter, provider) + svc.auth = &stubAuthService{authCtx: domain.AuthContext{ + Account: domain.Account{ID: "acc1", Status: domain.AccountStatusActive, IsGuest: isGuest}, + APIKey: domain.APIKey{ID: "key1", Active: true}, + }} + svc.SetGuestMaxOutputTokens(50) + svc.catalog.(*stubModelCatalog).model.GuestAllowed = true + return svc +} + +func proxyChat(t *testing.T, svc *GenerateService, body string) { + t.Helper() + req, err := mapper.ChatCompletionRequestToDomain([]byte(body)) + if err != nil { + t.Fatalf("map request: %v", err) + } + if _, err := svc.Proxy(context.Background(), domain.EndpointChatCompletions, []byte(body), req, "sk-test"); err != nil { + t.Fatalf("Proxy: %v", err) + } +} + +func sentField(t *testing.T, raw, field string) any { + t.Helper() + var body map[string]any + if err := json.Unmarshal([]byte(raw), &body); err != nil { + t.Fatalf("provider body: %v", err) + } + return body[field] +} + +// A guest's chat request reaches the provider with an output limit no larger +// than the guest limit (50 here), and Metering holds credit for that limit. +// Other accounts' requests keep their own limit, or none. +func TestProxy_GuestChatRequestsCarryTheOutputLimit(t *testing.T) { + const user = `"messages":[{"role":"user","content":"hi"}]` + for name, tc := range map[string]struct { + guest bool + body, field string + want float64 + }{ + "guest, no limit": {true, `{"model":"model",` + user + `}`, "max_completion_tokens", 50}, + "guest, larger limit": {true, `{"model":"model",` + user + `,"max_tokens":4000}`, "max_tokens", 50}, + "guest, smaller limit is kept": {true, `{"model":"model",` + user + `,"max_tokens":20}`, "max_tokens", 20}, + "guest, max_completion_tokens": {true, `{"model":"model",` + user + `,"max_completion_tokens":4000}`, "max_completion_tokens", 50}, + "account, limit kept": {false, `{"model":"model",` + user + `,"max_tokens":100000}`, "max_tokens", 100000}, + "account, small limit kept": {false, `{"model":"model",` + user + `,"max_tokens":80}`, "max_tokens", 80}, + } { + t.Run(name, func(t *testing.T) { + meter, provider := &stubUsageMeter{}, &stubProvider{} + proxyChat(t, newGuestTestService(meter, provider, tc.guest), tc.body) + if got := sentField(t, provider.lastRaw, tc.field); got != tc.want { + t.Fatalf("provider got %s = %v, want %v", tc.field, got, tc.want) + } + if got := meter.lastReserveReq.MaxOutputTokens; got == nil || float64(*got) != tc.want { + t.Fatalf("held output tokens = %v, want %v", got, tc.want) + } + }) + } +} + +// Apart from the output limit, a guest's request is an ordinary request: +// other fields reach the provider as sent. +func TestProxy_GuestRequestsAreOtherwiseOrdinary(t *testing.T) { + body := `{"model":"model","messages":[{"role":"user","content":"hi"}],"tools":[{"type":"function","function":{"name":"f","parameters":{"type":"object"}}}]}` + meter, provider := &stubUsageMeter{}, &stubProvider{} + proxyChat(t, newGuestTestService(meter, provider, true), body) + if sentField(t, provider.lastRaw, "tools") == nil { + t.Fatal("a guest's tools were dropped") + } + + // A signed-in request without a limit goes upstream without one, and + // Metering holds the model's maximum output. + meter, provider = &stubUsageMeter{}, &stubProvider{} + proxyChat(t, newGuestTestService(meter, provider, false), body) + if got := sentField(t, provider.lastRaw, "max_tokens"); got != nil { + t.Fatalf("a signed-in request got max_tokens %v", got) + } + if got := sentField(t, provider.lastRaw, "max_completion_tokens"); got != nil { + t.Fatalf("a signed-in request got max_completion_tokens %v", got) + } + if meter.lastReserveReq.MaxOutputTokens != nil { + t.Fatalf("a signed-in request's hold got an output limit: %d", *meter.lastReserveReq.MaxOutputTokens) + } +} diff --git a/apps/gateway/internal/application/services/proxy.go b/apps/gateway/internal/application/services/proxy.go index 5dde257..83a63d4 100644 --- a/apps/gateway/internal/application/services/proxy.go +++ b/apps/gateway/internal/application/services/proxy.go @@ -56,6 +56,18 @@ func (s *GenerateService) Proxy(ctx context.Context, endpoint string, raw []byte } return domain.Annotate(err, "account_id", authCtx.Account.ID, "model", model.PublicModelID, "provider", model.ProviderConfig.ProviderName) } + if authCtx.Account.IsGuest { + err := s.checkGuestModel(ctx, endpoint, req.PublicModelID) + if err == nil { + raw, req, err = capOutput(raw, req, s.guestMaxOutputTokens) + } + if err != nil { + err = clientGone(ctx, err) + s.recordTerminalOutcome(metrics.OutcomeError, endpoint, req, nil, time.Since(start).Milliseconds()) + s.recordError(endpoint, req, nil, metrics.PhaseRequest, err) + return nil, reject(err, nil) + } + } var model domain.PublicModel execReq := req if endpoint == domain.EndpointSystemOne { @@ -83,6 +95,10 @@ func (s *GenerateService) Proxy(ctx context.Context, endpoint string, raw []byte s.recordFailure(ctx, nil, &authCtx, endpoint, &execReq, &model, err, nil, time.Since(start).Milliseconds()) return nil, reject(err, &model) } + // The hold covers the prompt as sent upstream, after masking. + if tokens, ok := estimateInputTokens(endpoint, masked); ok { + execReq.EstimatedInputTokens = &tokens + } reservationID, err := s.metering.Reserve(ctx, authCtx, endpoint, execReq, model, uuid.NewString()) if err != nil { err = clientGone(ctx, err) @@ -90,6 +106,14 @@ func (s *GenerateService) Proxy(ctx context.Context, endpoint string, raw []byte s.recordError(endpoint, execReq, &model, metrics.PhaseRequest, err) return nil, reject(err, &model) } + // From here the provider runs and bills us, so the request no longer + // follows the client's connection. A client that hangs up still gets its + // answer finished, and is charged the usage the provider reports at the + // end. Cutting the answer off instead would leave no usage to charge, + // while the provider still bills what it generated. clientCtx only stops + // a fallback answer nobody would read. + clientCtx := ctx + ctx = context.WithoutCancel(ctx) ctx, cancelInference := context.WithCancelCause(ctx) stopRenewal := func() {} if keeper, ok := s.metering.(ports.ReservationKeeper); ok { @@ -164,7 +188,7 @@ func (s *GenerateService) Proxy(ctx context.Context, endpoint string, raw []byte call, err := attempt(model) executed := model - if err != nil && shouldTryFallback(ctx, err, model.Fallback) && + if err != nil && clientCtx.Err() == nil && shouldTryFallback(ctx, err, model.Fallback) && withProviderTarget(model, *model.Fallback).SupportsEndpoint(endpoint) { executed = withProviderTarget(model, *model.Fallback) s.logFallback(requestID, model, executed, sanitizeErrorWithPIIMapping(err, mapping)) @@ -609,9 +633,6 @@ func (p *ProxyStream) Next() ([]byte, error) { // Complete reports whether the provider finished the response cleanly. func (p *ProxyStream) Complete() bool { return p.sawDone || p.o.finish != nil } -// Finished reports whether a finish reason has passed. -func (p *ProxyStream) Finished() bool { return p.o.finish != nil } - // SawDone reports whether the provider sent [DONE]. func (p *ProxyStream) SawDone() bool { return p.sawDone } diff --git a/apps/gateway/internal/application/services/proxy_hangup_test.go b/apps/gateway/internal/application/services/proxy_hangup_test.go new file mode 100644 index 0000000..29e1a15 --- /dev/null +++ b/apps/gateway/internal/application/services/proxy_hangup_test.go @@ -0,0 +1,125 @@ +package services + +import ( + "context" + "io" + "testing" + + "github.com/dappnode/dappnode-nexus-gateway/apps/gateway/internal/application/ports" + "github.com/dappnode/dappnode-nexus-gateway/pkg/domain" +) + +// hangUpProvider records the context the gateway gives it. Its stream comes +// from a pipe the test writes to, closed like an HTTP body when that context +// is cancelled. +type hangUpProvider struct { + onCall func() // Called before the provider answers. + errAt error // The context's error after onCall. + ctx context.Context + body *io.PipeReader + write *io.PipeWriter +} + +func newHangUpProvider(onCall func()) *hangUpProvider { + pr, pw := io.Pipe() + return &hangUpProvider{onCall: onCall, body: pr, write: pw} +} + +func (p *hangUpProvider) call(ctx context.Context) { + p.ctx = ctx + if p.onCall != nil { + p.onCall() + } + p.errAt = ctx.Err() +} + +func (p *hangUpProvider) Complete(ctx context.Context, _ []byte, _ domain.PublicModel) ([]byte, *domain.TinfoilTransportProof, error) { + p.call(ctx) + return []byte(stubAnswer), nil, nil +} + +func (p *hangUpProvider) Stream(ctx context.Context, _ []byte, _ domain.PublicModel) (ports.ProviderStream, error) { + p.call(ctx) + context.AfterFunc(ctx, func() { _ = p.write.CloseWithError(ctx.Err()) }) + go func() { + _, _ = io.WriteString(p.write, "data: {\"id\":\"r\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"Hel\"}}]}\n\n") + }() + return ports.ProviderStream{Body: p.body}, nil +} + +func newHangUpService(meter *stubUsageMeter, provider ports.Provider) *GenerateService { + svc := newDirectModelGenerateService(meter, &stubProvider{}) + svc.registry = &stubProviderRegistry{providers: map[string]ports.Provider{"openai": provider}} + return svc +} + +func drain(stream *ProxyStream) { + for { + if _, err := stream.Next(); err != nil { + break + } + } + _ = stream.Close() +} + +// A client that hangs up mid-stream doesn't stop the provider: the answer +// runs to the end, and the usage it reports is charged. +func TestProxy_ClientThatHangsUpMidStreamIsCharged(t *testing.T) { + meter := &stubUsageMeter{} + client, hangUp := context.WithCancel(context.Background()) + provider := newHangUpProvider(nil) + svc := newHangUpService(meter, provider) + + call, err := svc.Proxy(client, domain.EndpointChatCompletions, []byte(`{"stream":true}`), domain.GenerateRequest{PublicModelID: "model", Stream: true}, "sk-test") + if err != nil { + t.Fatalf("Proxy: %v", err) + } + hangUp() + if err := provider.ctx.Err(); err != nil { + t.Fatalf("the client hanging up cancelled the provider stream: %v", err) + } + go func() { + _, _ = io.WriteString(provider.write, "data: {\"id\":\"r\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"lo\"},\"finish_reason\":\"stop\"}]}\n\n"+ + "data: {\"id\":\"r\",\"choices\":[],\"usage\":{\"prompt_tokens\":9,\"completion_tokens\":2,\"total_tokens\":11}}\n\n"+ + "data: [DONE]\n\n") + _ = provider.write.Close() + }() + drain(call.Stream) + if meter.successCalls != 1 || meter.failureCalls != 0 { + t.Fatalf("success %d, failure %d; want the finished answer charged once", meter.successCalls, meter.failureCalls) + } +} + +// The same holds without streaming. +func TestProxy_ClientThatHangsUpDuringACompleteAnswerIsCharged(t *testing.T) { + meter := &stubUsageMeter{} + client, hangUp := context.WithCancel(context.Background()) + provider := newHangUpProvider(hangUp) + svc := newHangUpService(meter, provider) + + if _, err := svc.Proxy(client, domain.EndpointChatCompletions, []byte(`{}`), domain.GenerateRequest{PublicModelID: "model"}, "sk-test"); err != nil { + t.Fatalf("Proxy: %v", err) + } + if provider.errAt != nil { + t.Fatalf("the client hanging up cancelled the provider call: %v", provider.errAt) + } + if meter.successCalls != 1 { + t.Fatalf("success %d, want the answer charged once", meter.successCalls) + } +} + +// No fallback answer is generated for a client that has already hung up. +func TestProxy_NoFallbackForAClientThatHungUp(t *testing.T) { + meter := &stubUsageMeter{} + client, hangUp := context.WithCancel(context.Background()) + primary := &stubProvider{beforeReturn: hangUp, err: domain.ErrProviderUnavailable("primary")} + fallback := &stubProvider{} + svc := newFallbackGenerateService(meter, primary, fallback) + + if _, err := svc.Proxy(client, domain.EndpointChatCompletions, []byte(`{}`), domain.GenerateRequest{PublicModelID: "model"}, "sk-test"); err == nil { + t.Fatal("expected the primary's error") + } + if fallback.calls != 0 { + t.Fatalf("fallback calls = %d, want 0", fallback.calls) + } +} diff --git a/apps/gateway/main.go b/apps/gateway/main.go index 182abf1..e1bc1a5 100644 --- a/apps/gateway/main.go +++ b/apps/gateway/main.go @@ -3,6 +3,7 @@ package main import ( "context" "crypto/rand" + "fmt" "log" "net/http" "os" @@ -93,6 +94,11 @@ func main() { generateSvc := services.NewGenerateService(authService, modelCatalog, routerClient, providerRegistry, meteringClient, piiFilter, zapLogger) generateSvc.SetPIIOptions(piiLang, piiFailOpen) generateSvc.SetTinfoilProofRepository(tinfoilProofRepo) + guestMaxOutput, err := positiveIntEnv("GUEST_MAX_OUTPUT_TOKENS", services.DefaultGuestMaxOutputTokens) + if err != nil { + log.Fatalf("invalid guest chat configuration: %v", err) + } + generateSvc.SetGuestMaxOutputTokens(guestMaxOutput) chatCompletionsSvc := services.NewChatCompletionsService(generateSvc, zapLogger) healthHandler := handlers.NewHealthHandler() @@ -247,3 +253,17 @@ func buildPIIFilter(log *logger.ZapLogger) (ports.PIIFilter, string, bool) { Logger: log, }), language, failOpen } + +// positiveIntEnv reads a positive whole number from the environment. A set but +// invalid value is an error, so a typo cannot silently loosen a limit. +func positiveIntEnv(name string, fallback int) (int, error) { + v := strings.TrimSpace(os.Getenv(name)) + if v == "" { + return fallback, nil + } + n, err := strconv.Atoi(v) + if err != nil || n <= 0 { + return 0, fmt.Errorf("%s must be a positive whole number", name) + } + return n, nil +} diff --git a/pkg/domain/account.go b/pkg/domain/account.go index 3a50b99..43fc994 100644 --- a/pkg/domain/account.go +++ b/pkg/domain/account.go @@ -30,6 +30,10 @@ type Account struct { // LastUsedAt is the most recent `last_used_at` across the account's API keys. // Populated only by list queries that aggregate key usage; nil elsewhere. LastUsedAt *time.Time + // IsGuest marks a guest chat account: a landing-page visitor without a + // login, funded with a small credit. Its chat requests get a smaller output + // limit than other accounts' (see services.capOutput). + IsGuest bool } // IsActive returns true if the account status is "active". diff --git a/pkg/domain/error.go b/pkg/domain/error.go index 303f03e..a8c8866 100644 --- a/pkg/domain/error.go +++ b/pkg/domain/error.go @@ -35,6 +35,7 @@ const ( ErrCodeUnknownField = "unknown_field" ErrCodeInsufficientBalance = "insufficient_balance" ErrCodeDailyTokenLimit = "daily_token_limit_exceeded" + ErrCodeGuestModelNotAllowed = "guest_model_not_allowed" ErrCodeConfidentialLimit = "confidential_request_limit_exceeded" ErrCodeClientCanceled = "client_canceled" ErrCodeInternalError = "internal_error" @@ -61,7 +62,7 @@ func (e *GatewayError) WithMeta(fields ...any) *GatewayError { if cp.Metadata == nil { cp.Metadata = make(map[string]any, len(fields)/2) } else { - m := make(map[string]any, len(cp.Metadata)+len(fields)/2) + m := make(map[string]any, len(cp.Metadata)) // Grows for the new fields. for k, v := range cp.Metadata { m[k] = v } @@ -93,6 +94,13 @@ func ErrInvalidAPIKey(msg string) *GatewayError { return &GatewayError{HTTPStatus: 401, Type: ErrTypeAuthentication, Code: ErrCodeInvalidAPIKey, Message: msg} } +// ErrGuestModelNotAllowed refuses a guest chat key a model or endpoint that +// is not open to guests (public_models.guest_allowed). +func ErrGuestModelNotAllowed(model string) *GatewayError { + return &GatewayError{HTTPStatus: 403, Type: ErrTypePermission, Code: ErrCodeGuestModelNotAllowed, + Message: fmt.Sprintf("Guest chat can't use '%s'. Sign up for free to use every model.", model)} +} + func ErrUnsupportedModel(model string) *GatewayError { return &GatewayError{HTTPStatus: 404, Type: ErrTypeInvalidRequest, Code: ErrCodeUnsupportedModel, Message: fmt.Sprintf("Model '%s' is not available", model)} } diff --git a/pkg/domain/generation.go b/pkg/domain/generation.go index 9daabf6..8a22524 100644 --- a/pkg/domain/generation.go +++ b/pkg/domain/generation.go @@ -27,10 +27,13 @@ type GenerateRequest struct { Input []InputItem Tools []ToolDefinition - Stream bool - MaxOutputTokens *int - ServiceTier *string - ParallelToolCalls *bool + Stream bool + MaxOutputTokens *int + // EstimatedInputTokens is an upper bound on the prompt's tokens, sized + // from the body sent upstream; nil when no bound is known. + EstimatedInputTokens *int64 + ServiceTier *string + ParallelToolCalls *bool // StructuredOutput is set when the client asks for JSON output. StructuredOutput bool } diff --git a/pkg/domain/model.go b/pkg/domain/model.go index 483e1bf..fd933ae 100644 --- a/pkg/domain/model.go +++ b/pkg/domain/model.go @@ -53,6 +53,9 @@ type PublicModel struct { CacheWritePricePerMillion *decimal.Decimal Currency string Active bool + // GuestAllowed marks a model guest chat keys may call; they can call no + // other model or router. + GuestAllowed bool } func (m PublicModel) EffectiveProofMode() string {