Request-sized holds, charge abandoned requests, guest output limit

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.

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)
+		}
+	}
+}

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 {

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) {

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)
 				}

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 {

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)

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 {

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

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 {

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
 }

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
+}