feat: keep the credit hold alive while inference runs

apps/gateway/internal/adapters/metering/client.go

@@ -11,6 +11,7 @@ import (
 	"reflect"
 	"strconv"
 	"strings"
+	"sync"
 	"time"
 
 	"github.com/dappnode/dappnode-nexus-gateway/apps/gateway/internal/adapters/observability/metrics"
@@ -19,9 +20,10 @@ import (
 )
 
 type Client struct {
-	baseURL string
-	token   string
-	http    *http.Client
+	baseURL      string
+	token        string
+	http         *http.Client
+	reservations sync.Map // Unclaimed leases by reservation ID, removed by KeepReservation.
 }
 
 func NewClient(baseURL, token string, timeout time.Duration) *Client {
@@ -58,12 +60,16 @@ func (c *Client) Reserve(ctx context.Context, auth domain.AuthContext, endpoint
 		MaxOutputTokens:        req.MaxOutputTokens,
 	}
 	var res reserveResponse
+	sent := time.Now() // A lease measured from here never outlasts Metering's.
 	if err := c.postJSON(ctx, "/internal/metering/reservations", payload, &res, 2); err != nil {
 		return "", err
 	}
 	if res.ReservationID == "" {
 		return "", meteringInternalError("reason", "missing reservation_id")
 	}
+	if res.LeaseMS > 0 {
+		c.reservations.Store(res.ReservationID, newLease(sent, res.LeaseMS))
+	}
 	return res.ReservationID, nil
 }
 
@@ -365,6 +371,8 @@ func meteringOperation(method, path string) string {
 		return "reserve"
 	case method == http.MethodPost && strings.Contains(requestPath, "/reservations/") && strings.HasSuffix(requestPath, "/complete"):
 		return "complete"
+	case method == http.MethodPost && strings.Contains(requestPath, "/reservations/") && strings.HasSuffix(requestPath, "/renew"):
+		return "renew"
 	case method == http.MethodPost && strings.HasSuffix(requestPath, "/usage-failures"):
 		return "record_failure"
 	case method == http.MethodPost && strings.HasSuffix(requestPath, "/authenticate"):
@@ -413,6 +421,7 @@ type reserveRequest struct {
 
 type reserveResponse struct {
 	ReservationID string `json:"reservation_id"`
+	LeaseMS       int64  `json:"lease_ms,omitempty"`
 }
 
 type completeRequest struct {

apps/gateway/internal/adapters/metering/reservation_lease.go

@@ -0,0 +1,88 @@
+package metering
+
+import (
+	"context"
+	"errors"
+	"time"
+
+	"github.com/dappnode/dappnode-nexus-gateway/pkg/domain"
+)
+
+// lease is a credit hold measured on this process's clock. Metering sends a
+// duration rather than a timestamp, so clock skew (an enclave) cannot make a
+// hold look expired or valid when it is not.
+type lease struct {
+	expires time.Time
+	length  time.Duration
+}
+
+func newLease(sent time.Time, ms int64) lease {
+	length := time.Duration(ms) * time.Millisecond
+	return lease{expires: sent.Add(length), length: length}
+}
+
+// KeepReservation renews a hold at a third of its remaining time until stop is
+// called. Transient failures retry within the lease, never more often than a
+// tenth of it or once per second. onLost runs when Metering says the hold is
+// gone, or when it expires without a successful renewal.
+func (c *Client) KeepReservation(parent context.Context, id string, onLost func()) func() {
+	value, ok := c.reservations.LoadAndDelete(id)
+	if !ok {
+		return func() {} // Older Metering responses do not advertise leases.
+	}
+	held := value.(lease)
+	floor := min(time.Second, held.length/10)
+	ctx, cancel := context.WithCancel(parent)
+	done := make(chan struct{})
+	lost := func() {
+		if ctx.Err() == nil {
+			onLost()
+		}
+	}
+	go func() {
+		defer close(done)
+		for {
+			remaining := time.Until(held.expires)
+			if remaining <= 0 {
+				lost()
+				return
+			}
+			timer := time.NewTimer(min(max(remaining/3, floor), remaining))
+			select {
+			case <-ctx.Done():
+				timer.Stop()
+				return
+			case <-timer.C:
+			}
+			if time.Until(held.expires) <= 0 {
+				continue
+			}
+			requestCtx, stop := context.WithDeadline(ctx, held.expires)
+			var result reserveResponse
+			sent := time.Now()
+			err := c.postJSON(requestCtx, "/internal/metering/reservations/"+id+"/renew", struct{}{}, &result, 0)
+			stop()
+			switch {
+			case err == nil && result.ReservationID == id && result.LeaseMS > 0:
+				if renewed := newLease(sent, result.LeaseMS); renewed.expires.After(held.expires) {
+					held = renewed
+				}
+			case rejectedByMetering(err):
+				lost() // Completed, expired, or unknown: retrying cannot restore it.
+				return
+			}
+		}
+	}()
+	return func() { cancel(); <-done }
+}
+
+// rejectedByMetering reports a 4xx answer; transport errors and 5xx are
+// transient and retried while the lease lasts.
+func rejectedByMetering(err error) bool {
+	var gwErr *domain.GatewayError
+	if !errors.As(err, &gwErr) {
+		return false
+	}
+	status, ok := gwErr.Metadata["metering_status"].(int)
+	return ok && status >= 400 && status < 500
+}

apps/gateway/internal/adapters/metering/reservation_lease_test.go

@@ -0,0 +1,137 @@
+package metering
+
+import (
+	"context"
+	"encoding/json"
+	"net/http"
+	"net/http/httptest"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"github.com/dappnode/dappnode-nexus-gateway/pkg/domain"
+)
+
+func TestReservationRenewsBeyondOriginalExpiryAndStops(t *testing.T) {
+	var calls atomic.Int32
+	var lost atomic.Bool
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if calls.Add(1) == 1 {
+			w.WriteHeader(503)
+			return
+		}
+		json.NewEncoder(w).Encode(reserveResponse{ReservationID: "lease", LeaseMS: 300})
+	}))
+	defer server.Close()
+	client := NewClient(server.URL, "fixture", time.Second)
+	client.reservations.Store("lease", newLease(time.Now(), 300))
+	stop := client.KeepReservation(context.Background(), "lease", func() { lost.Store(true) })
+	time.Sleep(600 * time.Millisecond)
+	stop()
+	if lost.Load() || calls.Load() < 3 {
+		t.Fatalf("renewal did not preserve hold: lost=%v calls=%d", lost.Load(), calls.Load())
+	}
+	before := calls.Load()
+	time.Sleep(150 * time.Millisecond)
+	if calls.Load() != before {
+		t.Fatal("renewal continued after stop")
+	}
+}
+
+func TestLostReservationStopsOnlyWhenLeaseExpires(t *testing.T) {
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.WriteHeader(503) }))
+	defer server.Close()
+	client := NewClient(server.URL, "fixture", time.Second)
+	client.reservations.Store("lease", newLease(time.Now(), 120))
+	lost := make(chan struct{})
+	stop := client.KeepReservation(context.Background(), "lease", func() { close(lost) })
+	defer stop()
+	select {
+	case <-lost:
+		t.Fatal("temporary error canceled a valid lease")
+	case <-time.After(50 * time.Millisecond):
+	}
+	select {
+	case <-lost:
+	case <-time.After(time.Second):
+		t.Fatal("lost hold did not stop inference")
+	}
+}
+
+// Metering refusing a renewal means the hold is gone; retrying cannot help.
+func TestRejectedRenewalIsLostImmediately(t *testing.T) {
+	var calls atomic.Int32
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		calls.Add(1)
+		w.WriteHeader(http.StatusConflict)
+		json.NewEncoder(w).Encode(map[string]any{"error": map[string]string{"code": "conflict", "message": "reservation is no longer active"}})
+	}))
+	defer server.Close()
+	client := NewClient(server.URL, "fixture", time.Second)
+	client.reservations.Store("lease", newLease(time.Now(), 3000))
+	lost := make(chan struct{})
+	stop := client.KeepReservation(context.Background(), "lease", func() { close(lost) })
+	defer stop()
+	select {
+	case <-lost:
+	case <-time.After(1500 * time.Millisecond):
+		t.Fatal("rejected renewal waited for the lease to run out")
+	}
+	if calls.Load() != 1 {
+		t.Fatalf("retried a rejected renewal: %d calls", calls.Load())
+	}
+}
+
+// During a Metering outage each request renews a bounded number of times
+// instead of an accelerating burst as the lease runs out.
+func TestRenewalRetriesAreSpacedDuringOutage(t *testing.T) {
+	var calls atomic.Int32
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		calls.Add(1)
+		w.WriteHeader(503)
+	}))
+	defer server.Close()
+	client := NewClient(server.URL, "fixture", time.Second)
+	client.reservations.Store("lease", newLease(time.Now(), 1000))
+	lost := make(chan struct{})
+	stop := client.KeepReservation(context.Background(), "lease", func() { close(lost) })
+	defer stop()
+	select {
+	case <-lost:
+	case <-time.After(3 * time.Second):
+		t.Fatal("lease never expired")
+	}
+	// One tenth of a 1 s lease is 100 ms, so at most about ten attempts.
+	if got := calls.Load(); got > 11 {
+		t.Fatalf("renewal attempts were not spaced: %d", got)
+	}
+}
+
+// The lease is measured on the Gateway's clock from before the request, so a
+// skewed Metering or enclave clock cannot shorten or extend it.
+func TestReserveStoresLeaseRelativeToGatewayClock(t *testing.T) {
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		json.NewEncoder(w).Encode(reserveResponse{ReservationID: "lease", LeaseMS: 60_000})
+	}))
+	defer server.Close()
+	client := NewClient(server.URL, "fixture", time.Second)
+	before := time.Now()
+	id, err := client.Reserve(context.Background(), domain.AuthContext{Account: domain.Account{ID: "a"}, APIKey: domain.APIKey{ID: "k"}}, "chat/completions", domain.GenerateRequest{}, domain.PublicModel{PublicModelID: "m"}, "request")
+	if err != nil || id != "lease" {
+		t.Fatalf("reserve = %q, %v", id, err)
+	}
+	value, ok := client.reservations.Load("lease")
+	if !ok {
+		t.Fatal("lease not stored")
+	}
+	held := value.(lease)
+	if held.length != time.Minute || held.expires.Before(before.Add(time.Minute)) || held.expires.After(time.Now().Add(time.Minute)) {
+		t.Fatalf("lease = %+v", held)
+	}
+}
+
+func TestRenewalMetricsDoNotIncludeReservationIDs(t *testing.T) {
+	if got := meteringOperation(http.MethodPost, "/internal/metering/reservations/private-id/renew"); got != "renew" {
+		t.Fatalf("renewal metric operation = %q", got)
+	}
+}

apps/gateway/internal/application/ports/usage.go

@@ -14,3 +14,9 @@ type UsageMeter interface {
 	RecordSuccess(ctx context.Context, reservationID string, auth domain.AuthContext, endpoint string, req domain.GenerateRequest, resp domain.GenerateResult, model domain.PublicModel, latencyMs int64) error
 	RecordFailure(ctx context.Context, reservationID *string, auth *domain.AuthContext, endpoint string, req *domain.GenerateRequest, model *domain.PublicModel, err error, partialUsage *domain.Usage, latencyMs int64) error
 }
+
+// ReservationKeeper maintains a live credit hold while inference progresses.
+// onLost stops inference only when renewal cannot preserve the hold.
+type ReservationKeeper interface {
+	KeepReservation(context.Context, string, func()) func()
+}

apps/gateway/internal/application/services/proxy.go

@@ -13,6 +13,7 @@ import (
 
 	"github.com/dappnode/dappnode-nexus-gateway/apps/gateway/internal/adapters/http/middleware"
 	"github.com/dappnode/dappnode-nexus-gateway/apps/gateway/internal/adapters/observability/metrics"
+	"github.com/dappnode/dappnode-nexus-gateway/apps/gateway/internal/application/ports"
 	"github.com/dappnode/dappnode-nexus-gateway/pkg/domain"
 	"github.com/dappnode/dappnode-nexus-gateway/pkg/observability/logfields"
 	"github.com/google/uuid"
@@ -72,7 +73,17 @@ func (s *GenerateService) Proxy(ctx context.Context, endpoint string, raw []byte
 		s.recordTerminalOutcome(metrics.OutcomeError, endpoint, execReq, &model, time.Since(start).Milliseconds())
 		return nil, err
 	}
+	ctx, cancelInference := context.WithCancelCause(ctx)
+	stopRenewal := func() {}
+	if keeper, ok := s.metering.(ports.ReservationKeeper); ok {
+		stopRenewal = keeper.KeepReservation(ctx, reservationID, func() { cancelInference(domain.ErrInternal("credit reservation could not be maintained")) })
+	}
+	stopLease := func() { stopRenewal(); cancelInference(nil) }
 	finish := func(target domain.PublicModel, o outcome) {
+		defer stopLease()
+		if cause, ok := context.Cause(ctx).(*domain.GatewayError); ok && o.err != nil {
+			o.err = cause
+		}
 		o.err = sanitizeErrorWithPIIMapping(o.err, mapping)
 		s.finishProxy(ctx, authCtx, endpoint, execReq, target, requestID, reservationID, start, o)
 	}
@@ -138,6 +149,10 @@ func (s *GenerateService) Proxy(ctx context.Context, endpoint string, raw []byte
 		call, err = attempt(executed)
 	}
 	if err != nil {
+		defer stopLease()
+		if cause, ok := context.Cause(ctx).(*domain.GatewayError); ok {
+			err = cause
+		}
 		err = sanitizeErrorWithPIIMapping(err, mapping)
 		fields := s.buildErrorLogFields(ctx, requestID, &authCtx, endpoint, execReq.PublicModelID, executed, err, time.Since(start).Milliseconds())
 		s.logger.Error("provider request failed", fields...)

apps/gateway/internal/application/services/reservation_keeper_test.go

@@ -0,0 +1,63 @@
+package services
+
+import (
+	"context"
+	"strings"
+	"testing"
+
+	"github.com/dappnode/dappnode-nexus-gateway/apps/gateway/internal/application/ports"
+	"github.com/dappnode/dappnode-nexus-gateway/pkg/domain"
+)
+
+// keeperMeter is a meter whose credit hold can be lost before inference ends.
+type keeperMeter struct {
+	stubUsageMeter
+	lose    bool
+	stopped bool
+}
+
+func (m *keeperMeter) KeepReservation(_ context.Context, _ string, onLost func()) func() {
+	if m.lose {
+		onLost()
+	}
+	return func() { m.stopped = true }
+}
+
+// waitingProvider answers only when its request is canceled, like a slow
+// upstream that is still working.
+type waitingProvider struct{ stubProvider }
+
+func (p *waitingProvider) Complete(ctx context.Context, _ []byte, _ domain.PublicModel) ([]byte, *domain.TinfoilTransportProof, error) {
+	<-ctx.Done()
+	return nil, nil, ctx.Err()
+}
+
+func keeperService(meter *keeperMeter, provider ports.Provider) *GenerateService {
+	return NewGenerateService(
+		&stubAuthService{authCtx: domain.AuthContext{Account: domain.Account{ID: "acc1", Status: domain.AccountStatusActive}, APIKey: domain.APIKey{ID: "key1", Active: true}}},
+		&stubModelCatalog{model: domain.PublicModel{PublicModelID: "model", ProviderModelID: "provider-model", ProviderConfig: domain.ProviderConfig{ProviderName: "openai"}, SupportsChatCompletions: true, MaxContextWindow: 1000, MaxOutputTokens: 100}},
+		nil, &stubProviderRegistry{providers: map[string]ports.Provider{"openai": provider}}, meter, nil, stubLogger{},
+	)
+}
+
+func TestProxy_StopsInferenceWhenCreditHoldIsLost(t *testing.T) {
+	meter := &keeperMeter{lose: true}
+	_, err := keeperService(meter, &waitingProvider{}).Proxy(context.Background(), domain.EndpointChatCompletions, []byte(`{}`), domain.GenerateRequest{PublicModelID: "model"}, "sk-test")
+	gwErr, ok := err.(*domain.GatewayError)
+	if !ok || gwErr.Code != domain.ErrCodeInternalError || !strings.Contains(gwErr.Message, "credit reservation") {
+		t.Fatalf("lost hold reported as %v, want an internal error, not a user cancellation", err)
+	}
+	if meter.failureCalls != 1 || !meter.stopped {
+		t.Fatalf("failure recorded %d times, keeper stopped %v", meter.failureCalls, meter.stopped)
+	}
+}
+
+func TestProxy_StopsKeepingHoldAfterCompletion(t *testing.T) {
+	meter := &keeperMeter{}
+	if _, err := keeperService(meter, &stubProvider{}).Proxy(context.Background(), domain.EndpointChatCompletions, []byte(`{}`), domain.GenerateRequest{PublicModelID: "model"}, "sk-test"); err != nil {
+		t.Fatal(err)
+	}
+	if meter.successCalls != 1 || !meter.stopped {
+		t.Fatalf("success recorded %d times, keeper stopped %v", meter.successCalls, meter.stopped)
+	}
+}