From f9a2953b4104aed21a8ace82b6e8023c55a35e7c Mon Sep 17 00:00:00 2001 From: Marketen Date: Wed, 4 Nov 2026 15:28:55 +0100 Subject: [PATCH] fix: charge requests whose client hangs up before the answer ends A request is charged when the provider reports its usage, which it does at the end of the answer. When the client hung up earlier, the gateway cut the provider off, no usage arrived, and Metering released the hold without a charge, while the provider still billed us for what it had generated. Anyone could get answers for free by disconnecting just before the end. Once credit is reserved, the provider call no longer follows the client's connection: the answer runs to the end, and its reported usage is charged, for every account. Pressing Stop therefore no longer lowers the bill; the request's max_tokens bounds it, and the provider client's existing idle timeout still ends a stalled stream. A client that hung up gets no fallback answer. The old 5-second post-finish drain is no longer needed. Co-Authored-By: Claude Opus 5.5 --- .../http/handlers/chat_completions_handler.go | 32 +---- .../handlers/chat_completions_stream_test.go | 4 +- .../http/handlers/stream_usage_repro_test.go | 25 ++-- .../internal/application/services/proxy.go | 13 +- .../application/services/proxy_hangup_test.go | 125 ++++++++++++++++++ 5 files changed, 149 insertions(+), 50 deletions(-) create mode 100644 apps/gateway/internal/application/services/proxy_hangup_test.go 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/application/services/proxy.go b/apps/gateway/internal/application/services/proxy.go index 5dde257..460a2a6 100644 --- a/apps/gateway/internal/application/services/proxy.go +++ b/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 } 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) + } +}