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..41867d1 100644 --- a/apps/gateway/internal/adapters/http/handlers/chat_completions_handler.go +++ b/apps/gateway/internal/adapters/http/handlers/chat_completions_handler.go @@ -94,7 +94,9 @@ func (h *ChatCompletionsHandler) relay(w http.ResponseWriter, r *http.Request, s writeMu.Lock() defer writeMu.Unlock() lastWrite = time.Now() - // A client that left keeps the loop draining the upstream for usage. + // A client that left fails the write; its canceled request context + // ends the upstream read on the next line, so there is nothing to do + // here. _ = sw.WriteLine(line) } stopKeepAlive := make(chan struct{}) 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..8e54b9f 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 @@ -1,6 +1,7 @@ package handlers import ( + "context" "errors" "io" "net/http/httptest" @@ -46,7 +47,7 @@ func TestRelay_EndsEveryStream(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(context.Background(), brokenBody{tc.body}, "m"), &finished) body := recorder.Body.String() for _, want := range tc.want { if !strings.Contains(body, want) { diff --git a/apps/gateway/internal/application/services/generate_service_test.go b/apps/gateway/internal/application/services/generate_service_test.go index b41aa8a..77f2307 100644 --- a/apps/gateway/internal/application/services/generate_service_test.go +++ b/apps/gateway/internal/application/services/generate_service_test.go @@ -139,8 +139,9 @@ func (s *stubProviderRegistry) GetProvider(name string) (ports.Provider, error) type stubProvider struct { beforeReturn func() err error - json string // Complete's body; a short answer when empty. - sse string // Stream's body. + json string // Complete's body; a short answer when empty. + sse string // Stream's body. + body io.ReadCloser // Stream's body when it has to fail mid-read. calls int lastRaw string } @@ -168,6 +169,9 @@ func (s *stubProvider) Stream(_ context.Context, raw []byte, _ domain.PublicMode if s.err != nil { return ports.ProviderStream{}, s.err } + if s.body != nil { + return ports.ProviderStream{Body: s.body}, nil + } return ports.ProviderStream{Body: io.NopCloser(strings.NewReader(s.sse))}, nil } diff --git a/apps/gateway/internal/application/services/proxy.go b/apps/gateway/internal/application/services/proxy.go index 5dde257..ff4a87e 100644 --- a/apps/gateway/internal/application/services/proxy.go +++ b/apps/gateway/internal/application/services/proxy.go @@ -148,7 +148,7 @@ func (s *GenerateService) Proxy(ctx context.Context, endpoint string, raw []byte if err != nil { return nil, err } - stream := newProxyStream(resp.Body, target.PublicModelID) + stream := newProxyStream(ctx, resp.Body, target.PublicModelID) stream.proof = resp.Proof stream.hideUsage = !wantsUsage(raw) stream.unmask = unmask(target) @@ -377,6 +377,9 @@ func observeJSON(body []byte, public string, unmask *unmasker) ([]byte, outcome, // the finish reason as they pass. Lines are forwarded exactly, except the // model field of JSON events. type ProxyStream struct { + // ctx is the request's context. A read that fails once it is canceled + // means the gateway cut the upstream, not that the provider failed. + ctx context.Context body io.ReadCloser scanner *bufio.Scanner public string @@ -411,16 +414,17 @@ func wantsUsage(raw []byte) bool { return json.Unmarshal(raw, &req) == nil && req.StreamOptions != nil && req.StreamOptions.IncludeUsage } -// NewProxyStream relays body, naming the public model. The service creates -// streams; it is exported for handler tests. -func NewProxyStream(body io.ReadCloser, public string) *ProxyStream { - return newProxyStream(body, public) +// NewProxyStream relays body, naming the public model on the request's +// context, which tells a caller hang-up from a provider cut. The service +// creates streams; it is exported for handler tests. +func NewProxyStream(ctx context.Context, body io.ReadCloser, public string) *ProxyStream { + return newProxyStream(ctx, body, public) } -func newProxyStream(body io.ReadCloser, public string) *ProxyStream { +func newProxyStream(ctx context.Context, body io.ReadCloser, public string) *ProxyStream { scanner := bufio.NewScanner(body) scanner.Buffer(make([]byte, 0, 64<<10), maxSSELine) - return &ProxyStream{body: body, scanner: scanner, public: public} + return &ProxyStream{ctx: ctx, body: body, scanner: scanner, public: public} } // prime reads until the first output, so an upstream that fails first can @@ -596,6 +600,11 @@ func (p *ProxyStream) Next() ([]byte, error) { p.end("eof_without_finish") case p.o.finish != nil: p.end("read_error_after_finish") // Complete; only trailing bytes were lost. + case p.ctx != nil && p.ctx.Err() != nil: + // The caller hung up, so the gateway cut the upstream read: the + // caller's cancellation, not a provider that failed. + p.o.err = domain.ErrClientCanceled() + p.end("client_gone") default: p.o.err = domain.ErrProviderError(http.StatusBadGateway, "provider stream failed") p.end("read_error") diff --git a/apps/gateway/internal/application/services/proxy_test.go b/apps/gateway/internal/application/services/proxy_test.go index 3793732..d8fa5e5 100644 --- a/apps/gateway/internal/application/services/proxy_test.go +++ b/apps/gateway/internal/application/services/proxy_test.go @@ -3,6 +3,7 @@ package services import ( "context" "encoding/json" + "errors" "io" "strings" "testing" @@ -173,6 +174,66 @@ func TestProxyStream_ClientLeaving(t *testing.T) { } } +// cutStream serves content, then fails the read the way a cut upstream does. +type cutStream struct { + content io.Reader + err error +} + +func (c *cutStream) Read(p []byte) (int, error) { + n, err := c.content.Read(p) + if err == io.EOF { + return n, c.err + } + return n, err +} + +func (c *cutStream) Close() error { return nil } + +// A read that breaks mid-answer is the caller's cancellation when the +// request context is already canceled (the caller hung up and the gateway +// cut the upstream), and the provider's failure only while it is alive. +func TestProxyStream_ReadErrorIsClassifiedByWhoLeft(t *testing.T) { + first := "data: {\"choices\":[{\"delta\":{\"content\":\"par\"}}]}\n\n" + for name, tc := range map[string]struct { + hungUp bool + fault string + code string + }{ + "caller hung up": {hungUp: true, fault: domain.FaultCaller, code: domain.ErrCodeClientCanceled}, + "provider cut": {hungUp: false, fault: domain.FaultProvider, code: domain.ErrCodeProviderError}, + } { + t.Run(name, func(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + meter := &resultMeter{} + primary := &proxyStub{body: &cutStream{content: strings.NewReader(first), err: errors.New("connection reset")}} + svc := proxyService(meter, "", map[string]ports.Provider{"primary": primary}, false) + call, err := svc.Proxy(ctx, domain.EndpointChatCompletions, []byte(`{}`), streamRequest(), "key") + if err != nil { + t.Fatal(err) + } + if tc.hungUp { + cancel() + } + if _, err := relayAll(t, call.Stream); err == nil { + t.Fatal("read error not reported") + } + call.Stream.Close() + var gwErr *domain.GatewayError + if meter.failed == nil || !errors.As(meter.failed, &gwErr) { + t.Fatalf("failure not metered: %v", meter.failed) + } + if domain.FaultOf(meter.failed) != tc.fault || gwErr.Code != tc.code { + t.Fatalf("failure = %s/%d %s, want %s", domain.FaultOf(meter.failed), gwErr.HTTPStatus, gwErr.Code, tc.fault) + } + if tc.hungUp && meter.successCalls != 0 { + t.Fatalf("abandoned stream billed as success") + } + }) + } +} + func TestProxyJSON(t *testing.T) { primary := &proxyStub{json: `{"id":"r1","model":"glm-5-3-flash","choices":[{"message":{"role":"assistant","content":"hi","reasoning_content":"why"},"finish_reason":"stop"}],"usage":{"prompt_tokens":2,"completion_tokens":1}}`} meter := &resultMeter{} diff --git a/apps/gateway/internal/application/services/reservation_keeper_test.go b/apps/gateway/internal/application/services/reservation_keeper_test.go index c38d864..84dd22f 100644 --- a/apps/gateway/internal/application/services/reservation_keeper_test.go +++ b/apps/gateway/internal/application/services/reservation_keeper_test.go @@ -2,6 +2,7 @@ package services import ( "context" + "errors" "strings" "testing" @@ -32,7 +33,7 @@ func (p *waitingProvider) Complete(ctx context.Context, _ []byte, _ domain.Publi return nil, nil, ctx.Err() } -func keeperService(meter *keeperMeter, provider ports.Provider) *GenerateService { +func keeperService(meter ports.UsageMeter, 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}}, @@ -61,3 +62,38 @@ func TestProxy_StopsKeepingHoldAfterCompletion(t *testing.T) { t.Fatalf("success recorded %d times, keeper stopped %v", meter.successCalls, meter.stopped) } } + +// streamKeeperMeter loses the hold while a stream is being read, and keeps +// the failure it was handed. +type streamKeeperMeter struct { + keeperMeter + failed error +} + +func (m *streamKeeperMeter) RecordFailure(ctx context.Context, id *string, a *domain.AuthContext, e string, req *domain.GenerateRequest, model *domain.PublicModel, err error, u *domain.Usage, l int64) error { + m.failed = err + return m.keeperMeter.RecordFailure(ctx, id, a, e, req, model, err, u, l) +} + +// A stream the gateway cuts because the credit hold was lost stays Nexus's +// failure: a canceled request context must not pass it off as a caller abort. +func TestProxyStream_LostCreditHoldOutranksACanceledRequest(t *testing.T) { + meter := &streamKeeperMeter{keeperMeter: keeperMeter{lose: true}} + provider := &stubProvider{body: &cutStream{ + content: strings.NewReader("data: {\"choices\":[{\"delta\":{\"content\":\"a\"}}]}\n\n"), + err: errors.New("connection reset"), + }} + call, err := keeperService(meter, provider).Proxy(context.Background(), domain.EndpointChatCompletions, []byte(`{}`), + domain.GenerateRequest{PublicModelID: "model", Stream: true}, "sk-test") + if err != nil { + t.Fatal(err) + } + if _, err := relayAll(t, call.Stream); err == nil { + t.Fatal("read error not reported") + } + call.Stream.Close() + gwErr, ok := meter.failed.(*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 caller cancellation", meter.failed) + } +}