Charge requests whose client hangs up before the answer ends

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/application/services/proxy.go

@@ -90,6 +90,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 +172,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 +617,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 }
 

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