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