diff --git a/README.md b/README.md index d6eb701..93243e1 100644 --- a/README.md +++ b/README.md @@ -138,3 +138,13 @@ 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 keys (accounts flagged `is_guest` by Metering) are narrowed by +`services.GuestPolicy`: only the models in `GUEST_MODELS` (default +`deepseek/deepseek-v4.1-flash`), requests up to `GUEST_MAX_REQUEST_BYTES` +(default 16384), output capped at `GUEST_MAX_OUTPUT_TOKENS` (default 1024) and +written into the request, and only plain chat fields with text content. The +reservation for a guest request is bounded by the bytes sent upstream rather +than the model's context window, so it fits a guest's small credit. Invalid +values stop the gateway at startup. The Nexus repository's +`docs/guest-chat.md` describes the whole design. diff --git a/apps/gateway/guest_config.go b/apps/gateway/guest_config.go new file mode 100644 index 0000000..72ef832 --- /dev/null +++ b/apps/gateway/guest_config.go @@ -0,0 +1,61 @@ +package main + +import ( + "fmt" + "strconv" + "strings" + + "github.com/dappnode/dappnode-nexus-gateway/apps/gateway/internal/application/services" +) + +// Guest chat defaults. At DeepSeek V4.1 Flash's prices the largest guest +// request (16 KiB in, 1,024 tokens out) reserves about 0.7 cents, inside the +// 1-cent credit the control plane grants each guest; shorter messages reserve +// proportionally less. Change these together with that credit; see the Nexus +// repo's docs/guest-chat.md. +const ( + defaultGuestModels = "deepseek/deepseek-v4.1-flash" + defaultGuestMaxRequestBytes = 16 << 10 + defaultGuestMaxOutputTokens = 1024 +) + +// guestPolicyFromEnv reads the guest chat limits. A set but invalid value is +// an error, so a typo cannot loosen them. +func guestPolicyFromEnv(getenv func(string) string) (services.GuestPolicy, error) { + policy := services.GuestPolicy{Models: map[string]bool{}} + models := defaultGuestModels + if v, set := lookup(getenv, "GUEST_MODELS"); set { + models = v + } + for _, model := range strings.Split(models, ",") { + if model = strings.TrimSpace(model); model != "" { + policy.Models[model] = true + } + } + + var err error + if policy.MaxRequestBytes, err = positiveIntEnv(getenv, "GUEST_MAX_REQUEST_BYTES", defaultGuestMaxRequestBytes); err != nil { + return services.GuestPolicy{}, err + } + if policy.MaxOutputTokens, err = positiveIntEnv(getenv, "GUEST_MAX_OUTPUT_TOKENS", defaultGuestMaxOutputTokens); err != nil { + return services.GuestPolicy{}, err + } + return policy, nil +} + +func lookup(getenv func(string) string, name string) (string, bool) { + v := strings.TrimSpace(getenv(name)) + return v, v != "" +} + +func positiveIntEnv(getenv func(string) string, name string, fallback int) (int, error) { + v, set := lookup(getenv, name) + if !set { + 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/apps/gateway/guest_config_test.go b/apps/gateway/guest_config_test.go new file mode 100644 index 0000000..62ce03b --- /dev/null +++ b/apps/gateway/guest_config_test.go @@ -0,0 +1,51 @@ +package main + +import "testing" + +func envFrom(values map[string]string) func(string) string { + return func(name string) string { return values[name] } +} + +func TestGuestPolicyFromEnv_Defaults(t *testing.T) { + policy, err := guestPolicyFromEnv(envFrom(nil)) + if err != nil { + t.Fatalf("guestPolicyFromEnv: %v", err) + } + if len(policy.Models) != 1 || !policy.Models[defaultGuestModels] { + t.Fatalf("models = %v", policy.Models) + } + if policy.MaxRequestBytes != 16384 || policy.MaxOutputTokens != 1024 { + t.Fatalf("limits = %d bytes, %d tokens", policy.MaxRequestBytes, policy.MaxOutputTokens) + } +} + +func TestGuestPolicyFromEnv_ReadsOverrides(t *testing.T) { + policy, err := guestPolicyFromEnv(envFrom(map[string]string{ + "GUEST_MODELS": " a/model , b/model ,", + "GUEST_MAX_REQUEST_BYTES": "8192", + "GUEST_MAX_OUTPUT_TOKENS": "512", + })) + if err != nil { + t.Fatalf("guestPolicyFromEnv: %v", err) + } + if len(policy.Models) != 2 || !policy.Models["a/model"] || !policy.Models["b/model"] { + t.Fatalf("models = %v", policy.Models) + } + if policy.MaxRequestBytes != 8192 || policy.MaxOutputTokens != 512 { + t.Fatalf("limits = %d bytes, %d tokens", policy.MaxRequestBytes, policy.MaxOutputTokens) + } +} + +// A typo must stop the gateway, not silently loosen or drop a guest limit. +func TestGuestPolicyFromEnv_RejectsInvalidLimits(t *testing.T) { + for _, env := range []map[string]string{ + {"GUEST_MAX_REQUEST_BYTES": "16k"}, + {"GUEST_MAX_REQUEST_BYTES": "0"}, + {"GUEST_MAX_OUTPUT_TOKENS": "-1"}, + {"GUEST_MAX_OUTPUT_TOKENS": "1e3"}, + } { + if _, err := guestPolicyFromEnv(envFrom(env)); err == nil { + t.Errorf("accepted %v", env) + } + } +} diff --git a/apps/gateway/internal/adapters/metering/client.go b/apps/gateway/internal/adapters/metering/client.go index f95852e..25caed6 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, + MaxInputTokens: req.MaxInputTokens, } 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"` + MaxInputTokens *int `json:"max_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..0f49129 100644 --- a/apps/gateway/internal/adapters/metering/client_test.go +++ b/apps/gateway/internal/adapters/metering/client_test.go @@ -282,3 +282,35 @@ func TestClientReserve_MapsDailyTokenLimitWithoutRetry(t *testing.T) { t.Fatalf("quota error retried %d times", attempts.Load()) } } + +// The prompt bound the gateway computed for a guest reaches metering, which +// reserves for it instead of the model's context window. +func TestClientReserve_SendsTheInputTokenBound(t *testing.T) { + var got map[string]any + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if err := json.NewDecoder(r.Body).Decode(&got); err != nil { + t.Errorf("decode: %v", err) + } + _, _ = w.Write([]byte(`{"reservation_id":"res-1"}`)) + })) + defer server.Close() + + bound, output := 1234, 50 + req := domain.GenerateRequest{PublicModelID: "m", MaxInputTokens: &bound, MaxOutputTokens: &output} + if _, err := NewClient(server.URL, "token", time.Second).Reserve(context.Background(), domain.AuthContext{}, "chat_completions", req, domain.PublicModel{PublicModelID: "m"}, "req-1"); err != nil { + t.Fatalf("Reserve: %v", err) + } + if got["max_input_tokens"] != float64(1234) || got["max_output_tokens"] != float64(50) { + t.Fatalf("reservation body = %v", got) + } + + // Without a bound the field is left out, so metering uses the context window. + got = nil + req.MaxInputTokens = nil + if _, err := NewClient(server.URL, "token", time.Second).Reserve(context.Background(), domain.AuthContext{}, "chat_completions", req, domain.PublicModel{PublicModelID: "m"}, "req-2"); err != nil { + t.Fatalf("Reserve: %v", err) + } + if _, present := got["max_input_tokens"]; present { + t.Fatalf("unbounded request sent max_input_tokens: %v", got) + } +} diff --git a/apps/gateway/internal/adapters/metering/runtime_client.go b/apps/gateway/internal/adapters/metering/runtime_client.go index ba39c2d..13df452 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, @@ -413,6 +414,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 { diff --git a/apps/gateway/internal/adapters/metering/runtime_client_test.go b/apps/gateway/internal/adapters/metering/runtime_client_test.go index 2ad1ca5..f3d5d94 100644 --- a/apps/gateway/internal/adapters/metering/runtime_client_test.go +++ b/apps/gateway/internal/adapters/metering/runtime_client_test.go @@ -52,6 +52,31 @@ func TestRuntimeClientAuthenticateHashesAPIKeyLocally(t *testing.T) { } } +// The guest flag decides whether the guest policy 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/application/services/generate_service.go b/apps/gateway/internal/application/services/generate_service.go index feeac0f..9e75048 100644 --- a/apps/gateway/internal/application/services/generate_service.go +++ b/apps/gateway/internal/application/services/generate_service.go @@ -25,6 +25,7 @@ type GenerateService struct { piiLang string piiFailOpen bool tinfoilProofs ports.TinfoilTransportProofRepository + guest *GuestPolicy logger ports.Logger } @@ -58,6 +59,20 @@ func (s *GenerateService) SetPIIOptions(language string, failOpen bool) { s.piiFailOpen = failOpen } +// SetGuestPolicy sets the limits for guest chat keys. Without a policy, +// guest keys are refused. +func (s *GenerateService) SetGuestPolicy(policy GuestPolicy) { + s.guest = &policy +} + +// applyGuestPolicy narrows a guest's request; see GuestPolicy. +func (s *GenerateService) applyGuestPolicy(endpoint string, raw []byte, req domain.GenerateRequest) ([]byte, domain.GenerateRequest, error) { + if s.guest == nil { + return nil, req, domain.ErrGuestNotAllowed("this gateway") + } + return s.guest.Apply(endpoint, raw, req) +} + func (s *GenerateService) SetTinfoilProofRepository(proofs ports.TinfoilTransportProofRepository) { s.tinfoilProofs = proofs } diff --git a/apps/gateway/internal/application/services/guest_policy.go b/apps/gateway/internal/application/services/guest_policy.go new file mode 100644 index 0000000..bc38979 --- /dev/null +++ b/apps/gateway/internal/application/services/guest_policy.go @@ -0,0 +1,142 @@ +package services + +import ( + "bytes" + "encoding/json" + "fmt" + + "github.com/dappnode/dappnode-nexus-gateway/pkg/domain" +) + +// GuestPolicy narrows what a guest chat key can do. A guest is a landing-page +// visitor without a login, holding a key funded with a fraction of a cent +// (see the Nexus repo's docs/guest-chat.md). +// +// Metering refuses any request whose worst-case cost the guest's credit +// cannot cover. This policy keeps that worst case small and honest: +// +// - the model is one of Models, never a router or a pricier model; +// - the request is at most MaxRequestBytes, so the prompt is at most that +// many tokens: tokenizers emit at most one token per byte of text, and the +// JSON around the text is longer than the chat template that replaces it; +// - the output is at most MaxOutputTokens, written into the request so the +// provider enforces it too; +// - only plain chat fields and text content are accepted, so no provider +// option (reasoning budgets, tools, images) can raise the cost. +type GuestPolicy struct { + Models map[string]bool + MaxRequestBytes int + MaxOutputTokens int +} + +// guestPromptOverheadTokens covers what reaches the model beyond the bytes +// the gateway reserved for: the chat template's role markers and the fields +// the provider adapter adds (model name, stream options). +const guestPromptOverheadTokens = 256 + +// guestRequestFields are the top-level request fields a guest may send. +var guestRequestFields = map[string]bool{ + "model": true, "messages": true, "stream": true, "stream_options": true, + "max_tokens": true, "max_completion_tokens": true, + "temperature": true, "top_p": true, "presence_penalty": true, "frequency_penalty": true, + "stop": true, "seed": true, "user": true, +} + +// guestMessageFields are the fields a guest message may carry. +var guestMessageFields = map[string]bool{"role": true, "content": true, "name": true} + +// guestRoles are the message roles a guest may send. +var guestRoles = map[string]bool{"system": true, "user": true, "assistant": true} + +// Apply checks a guest's request and returns the body to forward, with its +// output limit set, and the request summary updated to match. +func (p GuestPolicy) Apply(endpoint string, raw []byte, req domain.GenerateRequest) ([]byte, domain.GenerateRequest, error) { + if endpoint != domain.EndpointChatCompletions { + return nil, req, domain.ErrGuestNotAllowed("this endpoint") + } + if len(raw) > p.MaxRequestBytes { + return nil, req, domain.ErrGuestRequestTooLarge(p.MaxRequestBytes) + } + if !p.Models[req.PublicModelID] { + return nil, req, domain.ErrGuestNotAllowed(fmt.Sprintf("the model '%s'", req.PublicModelID)) + } + + 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") + } + for field := range body { + if !guestRequestFields[field] { + return nil, req, domain.ErrGuestNotAllowed(fmt.Sprintf("the field '%s'", field)) + } + } + messages, ok := body["messages"].([]any) + if !ok { + return nil, req, domain.ErrInvalidField("messages must be an array") + } + for i, message := range messages { + if err := checkGuestMessage(i, message); err != nil { + return nil, req, err + } + } + + limit := p.MaxOutputTokens + if req.MaxOutputTokens != nil && *req.MaxOutputTokens > 0 && *req.MaxOutputTokens < limit { + limit = *req.MaxOutputTokens + } + // Keep the client's field name; providers that need max_tokens get it + // from the provider adapter. + if _, ok := body["max_completion_tokens"]; ok { + body["max_completion_tokens"] = limit + } else { + body["max_tokens"] = limit + } + rewritten, err := json.Marshal(body) + if err != nil { + return nil, req, domain.ErrInternal("an internal error occurred").WithMeta("reason", "re-encode guest request", "error", err.Error()) + } + req.MaxOutputTokens = &limit + return rewritten, req, nil +} + +// PromptTokenBound is the most prompt tokens a guest request can reach the +// model with, given the exact body sent upstream (after PII masking). +func (p GuestPolicy) PromptTokenBound(upstreamBody []byte) int { + return len(upstreamBody) + guestPromptOverheadTokens +} + +// checkGuestMessage accepts a system, user or assistant message whose content +// is text: a string, or text parts. +func checkGuestMessage(i int, message any) error { + fields, ok := message.(map[string]any) + if !ok { + return domain.ErrInvalidField(fmt.Sprintf("message[%d] must be an object", i)) + } + for field := range fields { + if !guestMessageFields[field] { + return domain.ErrGuestNotAllowed(fmt.Sprintf("the message field '%s'", field)) + } + } + if role, _ := fields["role"].(string); !guestRoles[role] { + return domain.ErrGuestNotAllowed(fmt.Sprintf("the message role '%v'", fields["role"])) + } + switch content := fields["content"].(type) { + case string: + return nil + case []any: + for _, part := range content { + p, ok := part.(map[string]any) + if !ok || p["type"] != "text" { + return domain.ErrGuestNotAllowed("images, audio or files") + } + if _, ok := p["text"].(string); !ok || len(p) != 2 { + return domain.ErrGuestNotAllowed("this message content") + } + } + return nil + default: + return domain.ErrInvalidField(fmt.Sprintf("message[%d] content must be text", i)) + } +} diff --git a/apps/gateway/internal/application/services/guest_policy_test.go b/apps/gateway/internal/application/services/guest_policy_test.go new file mode 100644 index 0000000..53ba7e6 --- /dev/null +++ b/apps/gateway/internal/application/services/guest_policy_test.go @@ -0,0 +1,208 @@ +package services + +import ( + "context" + "encoding/json" + "errors" + "strings" + "testing" + + "github.com/dappnode/dappnode-nexus-gateway/apps/gateway/internal/adapters/http/mapper" + "github.com/dappnode/dappnode-nexus-gateway/pkg/domain" +) + +var testGuestPolicy = GuestPolicy{ + Models: map[string]bool{"guest-model": true}, + MaxRequestBytes: 2048, + MaxOutputTokens: 512, +} + +// applyGuest maps a body the way the chat handler does, then applies the +// guest policy to it. +func applyGuest(t *testing.T, policy GuestPolicy, endpoint, body string) ([]byte, domain.GenerateRequest, error) { + t.Helper() + req, err := mapper.ChatCompletionRequestToDomain([]byte(body)) + if err != nil { + t.Fatalf("map request: %v", err) + } + return policy.Apply(endpoint, []byte(body), req) +} + +func decodeBody(t *testing.T, raw []byte) map[string]any { + t.Helper() + var body map[string]any + if err := json.Unmarshal(raw, &body); err != nil { + t.Fatalf("decode rewritten body: %v", err) + } + return body +} + +func TestGuestPolicy_WritesTheOutputLimitIntoTheRequest(t *testing.T) { + cases := []struct { + name, body, field string + want int + }{ + {"no limit sent", `{"model":"guest-model","messages":[{"role":"user","content":"hi"}]}`, "max_tokens", 512}, + {"higher limit is clamped", `{"model":"guest-model","messages":[{"role":"user","content":"hi"}],"max_tokens":100000}`, "max_tokens", 512}, + {"lower limit is kept", `{"model":"guest-model","messages":[{"role":"user","content":"hi"}],"max_tokens":64}`, "max_tokens", 64}, + {"zero limit means the cap", `{"model":"guest-model","messages":[{"role":"user","content":"hi"}],"max_tokens":0}`, "max_tokens", 512}, + {"client's field name is kept", `{"model":"guest-model","messages":[{"role":"user","content":"hi"}],"max_completion_tokens":9000}`, "max_completion_tokens", 512}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + raw, req, err := applyGuest(t, testGuestPolicy, domain.EndpointChatCompletions, tc.body) + if err != nil { + t.Fatalf("Apply: %v", err) + } + body := decodeBody(t, raw) + if got, _ := body[tc.field].(float64); int(got) != tc.want { + t.Fatalf("%s in body = %v, want %d", tc.field, body[tc.field], tc.want) + } + if req.MaxOutputTokens == nil || *req.MaxOutputTokens != tc.want { + t.Fatalf("request summary limit = %v, want %d", req.MaxOutputTokens, tc.want) + } + if _, both := body["max_tokens"]; both && tc.field == "max_completion_tokens" { + t.Fatal("the body carries both output-limit fields") + } + }) + } +} + +func TestGuestPolicy_AcceptsPlainTextChat(t *testing.T) { + body := `{"model":"guest-model","stream":true,"stream_options":{"include_usage":true},"temperature":1, + "messages":[{"role":"system","content":"You are Nexus."},{"role":"user","content":[{"type":"text","text":"hi"}]}, + {"role":"assistant","content":"hello"},{"role":"user","content":"again"}]}` + if _, _, err := applyGuest(t, testGuestPolicy, domain.EndpointChatCompletions, body); err != nil { + t.Fatalf("plain text chat refused: %v", err) + } +} + +// Everything that could raise a guest request's cost, or reach a model or +// feature the reservation does not cover, is refused. +func TestGuestPolicy_RefusesWhatGuestsCannotUse(t *testing.T) { + user := `{"role":"user","content":"hi"}` + cases := map[string]struct { + endpoint, body, code string + }{ + "other model": {domain.EndpointChatCompletions, `{"model":"pricey-model","messages":[` + user + `]}`, domain.ErrCodeGuestNotAllowed}, + "router": {domain.EndpointChatCompletions, `{"model":"nexus/auto","messages":[` + user + `]}`, domain.ErrCodeGuestNotAllowed}, + "other endpoint": {domain.EndpointSystemOne, `{"model":"guest-model","messages":[` + user + `]}`, domain.ErrCodeGuestNotAllowed}, + "reasoning option": {domain.EndpointChatCompletions, `{"model":"guest-model","messages":[` + user + `],"reasoning_effort":"high"}`, domain.ErrCodeGuestNotAllowed}, + "provider thinking": {domain.EndpointChatCompletions, `{"model":"guest-model","messages":[` + user + `],"thinking":{"type":"enabled"}}`, domain.ErrCodeGuestNotAllowed}, + "provider extra body": {domain.EndpointChatCompletions, `{"model":"guest-model","messages":[` + user + `],"extra_body":{}}`, domain.ErrCodeGuestNotAllowed}, + "tools": {domain.EndpointChatCompletions, `{"model":"guest-model","messages":[` + user + `],"tools":[{"type":"function","function":{"name":"f"}}]}`, domain.ErrCodeGuestNotAllowed}, + "tool role": {domain.EndpointChatCompletions, `{"model":"guest-model","messages":[` + user + `,{"role":"tool","tool_call_id":"1","content":"x"}]}`, domain.ErrCodeGuestNotAllowed}, + "developer role": {domain.EndpointChatCompletions, `{"model":"guest-model","messages":[{"role":"developer","content":"x"}]}`, domain.ErrCodeGuestNotAllowed}, + "image": {domain.EndpointChatCompletions, `{"model":"guest-model","messages":[{"role":"user","content":[{"type":"image_url","image_url":{"url":"https://x"}}]}]}`, domain.ErrCodeGuestNotAllowed}, + "text part extra": {domain.EndpointChatCompletions, `{"model":"guest-model","messages":[{"role":"user","content":[{"type":"text","text":"x","cache_control":{}}]}]}`, domain.ErrCodeGuestNotAllowed}, + "message extra field": {domain.EndpointChatCompletions, `{"model":"guest-model","messages":[{"role":"user","content":"x","reasoning_content":"y"}]}`, domain.ErrCodeGuestNotAllowed}, + "conversation too big": {domain.EndpointChatCompletions, `{"model":"guest-model","messages":[{"role":"user","content":"` + strings.Repeat("a", 2048) + `"}]}`, domain.ErrCodeGuestRequestTooLarge}, + } + for name, tc := range cases { + t.Run(name, func(t *testing.T) { + _, _, err := applyGuest(t, testGuestPolicy, tc.endpoint, tc.body) + var gwErr *domain.GatewayError + if !errors.As(err, &gwErr) || gwErr.Code != tc.code { + t.Fatalf("error = %v, want %s", err, tc.code) + } + }) + } +} + +func TestGuestPolicy_PromptTokenBoundCoversTheUpstreamBody(t *testing.T) { + body := []byte(`{"model":"guest-model","messages":[{"role":"user","content":"hi"}]}`) + if got := testGuestPolicy.PromptTokenBound(body); got != len(body)+guestPromptOverheadTokens { + t.Fatalf("bound = %d, want %d", got, len(body)+guestPromptOverheadTokens) + } +} + +func newGuestGenerateService(meter *stubUsageMeter, provider *stubProvider, isGuest bool) *GenerateService { + svc := newDirectModelGenerateService(meter, provider) + svc.auth = &stubAuthService{authCtx: domain.AuthContext{ + Account: domain.Account{ID: "guest-account", Status: domain.AccountStatusActive, IsGuest: isGuest}, + APIKey: domain.APIKey{ID: "guest-key", Active: true}, + }} + return svc +} + +func proxyGuestRequest(t *testing.T, svc *GenerateService, body string) error { + t.Helper() + req, err := mapper.ChatCompletionRequestToDomain([]byte(body)) + if err != nil { + t.Fatalf("map request: %v", err) + } + _, err = svc.Proxy(context.Background(), domain.EndpointChatCompletions, []byte(body), req, "sk-guest") + return err +} + +// A guest's reservation covers the prompt actually sent and the capped +// output, and the provider receives the capped request. +func TestProxy_GuestReservationIsBoundedByTheRequest(t *testing.T) { + meter, provider := &stubUsageMeter{}, &stubProvider{} + svc := newGuestGenerateService(meter, provider, true) + svc.SetGuestPolicy(GuestPolicy{Models: map[string]bool{"model": true}, MaxRequestBytes: 4096, MaxOutputTokens: 50}) + + if err := proxyGuestRequest(t, svc, `{"model":"model","messages":[{"role":"user","content":"hello"}],"max_tokens":1000}`); err != nil { + t.Fatalf("Proxy: %v", err) + } + got := meter.lastReserveReq + if got.MaxOutputTokens == nil || *got.MaxOutputTokens != 50 { + t.Fatalf("reserved output tokens = %v, want 50", got.MaxOutputTokens) + } + sent := provider.lastRaw + if got.MaxInputTokens == nil || *got.MaxInputTokens != len(sent)+guestPromptOverheadTokens { + t.Fatalf("reserved input tokens = %v, want %d (upstream body %d bytes + overhead)", got.MaxInputTokens, len(sent)+guestPromptOverheadTokens, len(sent)) + } + if body := decodeBody(t, []byte(sent)); body["max_tokens"] != float64(50) { + t.Fatalf("provider got max_tokens %v, want 50", body["max_tokens"]) + } +} + +// A refused guest request never reserves credit or reaches a provider. +func TestProxy_RefusedGuestRequestReachesNothing(t *testing.T) { + meter, provider := &stubUsageMeter{}, &stubProvider{} + svc := newGuestGenerateService(meter, provider, true) + svc.SetGuestPolicy(GuestPolicy{Models: map[string]bool{"model": true}, MaxRequestBytes: 4096, MaxOutputTokens: 50}) + + err := proxyGuestRequest(t, svc, `{"model":"model","messages":[{"role":"user","content":"hi"}],"thinking":{"type":"enabled"}}`) + var gwErr *domain.GatewayError + if !errors.As(err, &gwErr) || gwErr.Code != domain.ErrCodeGuestNotAllowed { + t.Fatalf("error = %v, want %s", err, domain.ErrCodeGuestNotAllowed) + } + if meter.reserveCalls != 0 || provider.calls != 0 { + t.Fatalf("refused guest request reserved %d times and called the provider %d times", meter.reserveCalls, provider.calls) + } +} + +// Without a configured policy a guest key is refused, never served unbounded. +func TestProxy_GuestWithoutPolicyIsRefused(t *testing.T) { + meter, provider := &stubUsageMeter{}, &stubProvider{} + svc := newGuestGenerateService(meter, provider, true) + + err := proxyGuestRequest(t, svc, `{"model":"model","messages":[{"role":"user","content":"hi"}]}`) + var gwErr *domain.GatewayError + if !errors.As(err, &gwErr) || gwErr.Code != domain.ErrCodeGuestNotAllowed { + t.Fatalf("error = %v, want %s", err, domain.ErrCodeGuestNotAllowed) + } + if meter.reserveCalls != 0 || provider.calls != 0 { + t.Fatal("a guest key was served without a guest policy") + } +} + +// Signed-in accounts are untouched by the guest policy. +func TestProxy_GuestPolicyLeavesOtherAccountsAlone(t *testing.T) { + meter, provider := &stubUsageMeter{}, &stubProvider{} + svc := newGuestGenerateService(meter, provider, false) + svc.SetGuestPolicy(GuestPolicy{Models: map[string]bool{"other": true}, MaxRequestBytes: 10, MaxOutputTokens: 1}) + + body := `{"model":"model","messages":[{"role":"user","content":"a request longer than ten bytes"}],"max_tokens":80,"reasoning_effort":"high"}` + if err := proxyGuestRequest(t, svc, body); err != nil { + t.Fatalf("Proxy: %v", err) + } + if meter.lastReserveReq.MaxInputTokens != nil { + t.Fatalf("a signed-in request got an input bound: %d", *meter.lastReserveReq.MaxInputTokens) + } + if got := meter.lastReserveReq.MaxOutputTokens; got == nil || *got != 80 { + t.Fatalf("a signed-in request's output limit changed: %v", got) + } +} diff --git a/apps/gateway/internal/application/services/proxy.go b/apps/gateway/internal/application/services/proxy.go index 5dde257..35cd529 100644 --- a/apps/gateway/internal/application/services/proxy.go +++ b/apps/gateway/internal/application/services/proxy.go @@ -56,6 +56,14 @@ 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 { + raw, req, err = s.applyGuestPolicy(endpoint, raw, req) + if err != nil { + 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 +91,12 @@ 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) } + if authCtx.Account.IsGuest { + // Reserve for the prompt actually sent, not the model's whole context + // window, so a short guest message fits in a guest's credit. + bound := s.guest.PromptTokenBound(masked) + execReq.MaxInputTokens = &bound + } reservationID, err := s.metering.Reserve(ctx, authCtx, endpoint, execReq, model, uuid.NewString()) if err != nil { err = clientGone(ctx, err) diff --git a/apps/gateway/main.go b/apps/gateway/main.go index 182abf1..912d192 100644 --- a/apps/gateway/main.go +++ b/apps/gateway/main.go @@ -4,9 +4,11 @@ import ( "context" "crypto/rand" "log" + "maps" "net/http" "os" "os/signal" + "slices" "strconv" "strings" "syscall" @@ -93,6 +95,13 @@ func main() { generateSvc := services.NewGenerateService(authService, modelCatalog, routerClient, providerRegistry, meteringClient, piiFilter, zapLogger) generateSvc.SetPIIOptions(piiLang, piiFailOpen) generateSvc.SetTinfoilProofRepository(tinfoilProofRepo) + guestPolicy, err := guestPolicyFromEnv(os.Getenv) + if err != nil { + log.Fatalf("invalid guest chat configuration: %v", err) + } + generateSvc.SetGuestPolicy(guestPolicy) + zapLogger.Info("guest chat limits", "models", strings.Join(slices.Sorted(maps.Keys(guestPolicy.Models)), ","), + "max_request_bytes", guestPolicy.MaxRequestBytes, "max_output_tokens", guestPolicy.MaxOutputTokens) chatCompletionsSvc := services.NewChatCompletionsService(generateSvc, zapLogger) healthHandler := handlers.NewHealthHandler() diff --git a/pkg/domain/account.go b/pkg/domain/account.go index 3a50b99..d85af3e 100644 --- a/pkg/domain/account.go +++ b/pkg/domain/account.go @@ -30,6 +30,9 @@ 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. Its requests are narrowed by services.GuestPolicy. + 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..23c71b8 100644 --- a/pkg/domain/error.go +++ b/pkg/domain/error.go @@ -36,6 +36,8 @@ const ( ErrCodeInsufficientBalance = "insufficient_balance" ErrCodeDailyTokenLimit = "daily_token_limit_exceeded" ErrCodeConfidentialLimit = "confidential_request_limit_exceeded" + ErrCodeGuestNotAllowed = "guest_not_allowed" + ErrCodeGuestRequestTooLarge = "guest_request_too_large" ErrCodeClientCanceled = "client_canceled" ErrCodeInternalError = "internal_error" ) @@ -147,6 +149,21 @@ func ErrConflict(msg string) *GatewayError { func ErrNotFound(resource, id string) *GatewayError { return &GatewayError{HTTPStatus: 404, Type: ErrTypeInvalidRequest, Code: ErrCodeNotFound, Message: fmt.Sprintf("%s '%s' not found", resource, id)} } + +// ErrGuestNotAllowed rejects a guest chat request that uses something guests +// cannot: another endpoint or model, a request field, or non-text content. +func ErrGuestNotAllowed(what string) *GatewayError { + return &GatewayError{HTTPStatus: 403, Type: ErrTypePermission, Code: ErrCodeGuestNotAllowed, + Message: fmt.Sprintf("Guest chat does not support %s. Sign up to use it.", what)} +} + +// ErrGuestRequestTooLarge rejects a guest chat request above the size limit: +// in practice a conversation that grew too long. +func ErrGuestRequestTooLarge(limitBytes int) *GatewayError { + return &GatewayError{HTTPStatus: 413, Type: ErrTypeInvalidRequest, Code: ErrCodeGuestRequestTooLarge, + Message: fmt.Sprintf("This conversation is too long for guest chat (limit %d bytes). Start a new chat or sign up to continue.", limitBytes)} +} + func ErrInsufficientBalance() *GatewayError { return &GatewayError{HTTPStatus: 402, Type: ErrTypePermission, Code: ErrCodeInsufficientBalance, Message: "Account has no spendable balance"} } diff --git a/pkg/domain/error_class.go b/pkg/domain/error_class.go index fab1a3a..6ba64c2 100644 --- a/pkg/domain/error_class.go +++ b/pkg/domain/error_class.go @@ -38,6 +38,8 @@ const ( var faultByCode = map[string]string{ ErrCodeInvalidAPIKey: FaultCaller, ErrCodeConfidentialLimit: FaultCaller, + ErrCodeGuestNotAllowed: FaultCaller, + ErrCodeGuestRequestTooLarge: FaultCaller, ErrCodeInactiveAPIKey: FaultCaller, ErrCodeInactiveAccount: FaultCaller, ErrCodeUnsupportedModel: FaultCaller, diff --git a/pkg/domain/generation.go b/pkg/domain/generation.go index 9daabf6..025b7b6 100644 --- a/pkg/domain/generation.go +++ b/pkg/domain/generation.go @@ -33,6 +33,10 @@ type GenerateRequest struct { ParallelToolCalls *bool // StructuredOutput is set when the client asks for JSON output. StructuredOutput bool + // MaxInputTokens bounds the prompt in metering's reservation when the + // gateway capped the request size (guest keys); nil means the model's + // context window. + MaxInputTokens *int } // InputItem is one message: its role and text.