diff --git a/backend/internal/handler/openai_gateway_handler.go b/backend/internal/handler/openai_gateway_handler.go index 83c4b4ac2..af4746531 100644 --- a/backend/internal/handler/openai_gateway_handler.go +++ b/backend/internal/handler/openai_gateway_handler.go @@ -1360,6 +1360,28 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) { reqLog.Info("openai.websocket_ingress_started") clientIP := ip.GetClientIP(c) userAgent := strings.TrimSpace(c.GetHeader("User-Agent")) + ctx := c.Request.Context() + maxIngressConnections := 0 + if h.cfg != nil { + maxIngressConnections = h.cfg.Gateway.OpenAIWS.MaxIngressConnectionsPerAPIKey + } + ingressLease, ingressLeaseAcquired, ingressLeaseErr := h.concurrencyHelper.AcquireOpenAIWSIngressLease(ctx, apiKey.ID, maxIngressConnections) + if ingressLeaseErr != nil { + reqLog.Error("openai.websocket_ingress_lease_acquire_failed", zap.Error(ingressLeaseErr)) + h.errorResponse(c, http.StatusServiceUnavailable, "service_unavailable", "WebSocket ingress capacity is temporarily unavailable") + return + } + if !ingressLeaseAcquired { + reqLog.Info("openai.websocket_ingress_capacity_rejected", zap.Int("max_ingress_connections_per_api_key", maxIngressConnections)) + c.Header("Retry-After", "5") + h.errorResponse(c, http.StatusTooManyRequests, "rate_limit_error", "Too many open WebSocket connections, please retry later") + return + } + if ingressLease != nil { + defer ingressLease.Release() + ctx = ingressLease.Context() + c.Request = c.Request.WithContext(ctx) + } wsConn, err := coderws.Accept(c.Writer, c.Request, &coderws.AcceptOptions{ CompressionMode: coderws.CompressionContextTakeover, @@ -1381,32 +1403,14 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) { }() wsConn.SetReadLimit(service.ResolveOpenAIWSClientReadLimitBytes(h.cfg)) - ctx := c.Request.Context() - maxIngressConnections := 0 - if h.cfg != nil { - maxIngressConnections = h.cfg.Gateway.OpenAIWS.MaxIngressConnectionsPerAPIKey - } - ingressLease, ingressLeaseAcquired, ingressLeaseErr := h.concurrencyHelper.AcquireOpenAIWSIngressLease(ctx, apiKey.ID, maxIngressConnections) - if ingressLeaseErr != nil { - reqLog.Error("openai.websocket_ingress_lease_acquire_failed", zap.Error(ingressLeaseErr)) - closeOpenAIClientWS(wsConn, coderws.StatusInternalError, "failed to reserve websocket ingress capacity") - return - } - if !ingressLeaseAcquired { - reqLog.Info("openai.websocket_ingress_capacity_rejected", zap.Int("max_ingress_connections_per_api_key", maxIngressConnections)) - closeOpenAIClientWS(wsConn, coderws.StatusTryAgainLater, "too many open websocket connections, please retry later") - return - } - if ingressLease != nil { - defer ingressLease.Release() - ctx = ingressLease.Context() - c.Request = c.Request.WithContext(ctx) - } - firstMessageTimeout := service.ResolveOpenAIWSClientFirstMessageTimeout(h.cfg) - readCtx, cancel := context.WithTimeout(ctx, firstMessageTimeout) - msgType, firstMessage, err := wsConn.Read(readCtx) - cancel() + msgType, firstMessage, err := service.ReadOpenAIWSClientMessage( + ctx, + wsConn, + firstMessageTimeout, + coderws.StatusPolicyViolation, + "missing first response.create message", + ) if err != nil { if errors.Is(context.Cause(ctx), service.ErrOpenAIWSIngressLeaseLost) { reqLog.Warn("openai.websocket_ingress_lease_lost_before_first_message", zap.Error(err)) diff --git a/backend/internal/handler/openai_gateway_handler_test.go b/backend/internal/handler/openai_gateway_handler_test.go index 66bffe5a8..851906122 100644 --- a/backend/internal/handler/openai_gateway_handler_test.go +++ b/backend/internal/handler/openai_gateway_handler_test.go @@ -857,17 +857,36 @@ func TestOpenAIResponsesWebSocket_IngressCapacityRejected(t *testing.T) { defer wsServer.Close() dialCtx, cancelDial := context.WithTimeout(context.Background(), 3*time.Second) - clientConn, _, err := coderws.Dial(dialCtx, "ws"+strings.TrimPrefix(wsServer.URL, "http")+"/openai/v1/responses", nil) + clientConn, response, err := coderws.Dial(dialCtx, "ws"+strings.TrimPrefix(wsServer.URL, "http")+"/openai/v1/responses", nil) cancelDial() - require.NoError(t, err) - defer func() { _ = clientConn.CloseNow() }() + require.Error(t, err) + require.Nil(t, clientConn) + require.NotNil(t, response) + require.Equal(t, http.StatusTooManyRequests, response.StatusCode) + _ = response.Body.Close() +} - readCtx, cancelRead := context.WithTimeout(context.Background(), 3*time.Second) - _, _, err = clientConn.Read(readCtx) - cancelRead() - var closeErr coderws.CloseError - require.ErrorAs(t, err, &closeErr) - require.Equal(t, coderws.StatusTryAgainLater, closeErr.Code) +func TestOpenAIResponsesWebSocket_IngressLeaseBackendUnavailableBeforeUpgrade(t *testing.T) { + gin.SetMode(gin.TestMode) + cache := &concurrencyCacheMock{ + acquireIngressLeaseFn: func(context.Context, int64, int, string) (bool, error) { + return false, errors.New("redis unavailable") + }, + } + h := newOpenAIHandlerForPreviousResponseIDValidation(t, cache) + h.cfg = &config.Config{} + h.cfg.Gateway.OpenAIWS.MaxIngressConnectionsPerAPIKey = 1 + wsServer := newOpenAIWSHandlerTestServer(t, h, middleware.AuthSubject{UserID: 1, Concurrency: 1}) + defer wsServer.Close() + + dialCtx, cancelDial := context.WithTimeout(context.Background(), 3*time.Second) + clientConn, response, err := coderws.Dial(dialCtx, "ws"+strings.TrimPrefix(wsServer.URL, "http")+"/openai/v1/responses", nil) + cancelDial() + require.Error(t, err) + require.Nil(t, clientConn) + require.NotNil(t, response) + require.Equal(t, http.StatusServiceUnavailable, response.StatusCode) + _ = response.Body.Close() } func TestOpenAIResponsesWebSocket_FirstMessageTimeoutUsesConfig(t *testing.T) { @@ -939,6 +958,33 @@ func TestOpenAIResponsesWebSocket_IngressLeaseReleasedOnEarlyReturn(t *testing.T }, time.Second, 10*time.Millisecond) } +func TestOpenAIResponsesWebSocket_IngressLeaseReleasedWhenUpgradeFails(t *testing.T) { + gin.SetMode(gin.TestMode) + cache := &concurrencyCacheMock{ + acquireIngressLeaseFn: func(context.Context, int64, int, string) (bool, error) { + return true, nil + }, + } + h := newOpenAIHandlerForPreviousResponseIDValidation(t, cache) + h.cfg = &config.Config{} + h.cfg.Gateway.OpenAIWS.MaxIngressConnectionsPerAPIKey = 1 + wsServer := newOpenAIWSHandlerTestServer(t, h, middleware.AuthSubject{UserID: 1, Concurrency: 1}) + defer wsServer.Close() + + req, err := http.NewRequest(http.MethodGet, wsServer.URL+"/openai/v1/responses", nil) + require.NoError(t, err) + req.Header.Set("Upgrade", "websocket") + req.Header.Set("Connection", "Upgrade") + req.Header.Set("Sec-WebSocket-Version", "13") + resp, err := http.DefaultClient.Do(req) + require.NoError(t, err) + _ = resp.Body.Close() + require.NotEqual(t, http.StatusSwitchingProtocols, resp.StatusCode) + require.Eventually(t, func() bool { + return atomic.LoadInt32(&cache.releaseIngressCalled) == 1 + }, time.Second, 10*time.Millisecond) +} + func TestOpenAIResponsesWebSocket_RejectsMessageIDAsPreviousResponseID(t *testing.T) { gin.SetMode(gin.TestMode) diff --git a/backend/internal/service/openai_ws_client_read.go b/backend/internal/service/openai_ws_client_read.go new file mode 100644 index 000000000..da8d650be --- /dev/null +++ b/backend/internal/service/openai_ws_client_read.go @@ -0,0 +1,70 @@ +package service + +import ( + "context" + "errors" + "time" + + coderws "github.com/coder/websocket" +) + +type openAIWSClientReadResult struct { + messageType coderws.MessageType + payload []byte + err error +} + +// ReadOpenAIWSClientMessage keeps one reader alive while control events send +// their close frame, then closes the transport and joins that reader. +func ReadOpenAIWSClientMessage( + controlCtx context.Context, + conn *coderws.Conn, + timeout time.Duration, + timeoutStatus coderws.StatusCode, + timeoutReason string, +) (coderws.MessageType, []byte, error) { + if conn == nil { + return 0, nil, errors.New("openai websocket client connection is nil") + } + if controlCtx == nil { + controlCtx = context.Background() + } + + readDone := make(chan openAIWSClientReadResult, 1) + go func() { + messageType, payload, err := conn.Read(context.Background()) + readDone <- openAIWSClientReadResult{messageType: messageType, payload: payload, err: err} + }() + + var timeoutCh <-chan time.Time + var timer *time.Timer + if timeout > 0 { + timer = time.NewTimer(timeout) + timeoutCh = timer.C + defer timer.Stop() + } + + closeAndJoin := func(status coderws.StatusCode, reason string, cause error) (coderws.MessageType, []byte, error) { + _ = conn.Close(status, reason) + _ = conn.CloseNow() + <-readDone + return 0, nil, NewOpenAIWSClientCloseError(status, reason, cause) + } + + select { + case result := <-readDone: + return result.messageType, result.payload, result.err + case <-timeoutCh: + return closeAndJoin(timeoutStatus, timeoutReason, context.DeadlineExceeded) + case <-controlCtx.Done(): + cause := context.Cause(controlCtx) + if errors.Is(cause, ErrOpenAIWSIngressLeaseLost) { + return closeAndJoin( + coderws.StatusTryAgainLater, + "websocket ingress capacity lease lost; please reconnect", + cause, + ) + } + return closeAndJoin(coderws.StatusGoingAway, "websocket request canceled", cause) + } +} diff --git a/backend/internal/service/openai_ws_client_read_test.go b/backend/internal/service/openai_ws_client_read_test.go new file mode 100644 index 000000000..82df19849 --- /dev/null +++ b/backend/internal/service/openai_ws_client_read_test.go @@ -0,0 +1,143 @@ +package service + +import ( + "context" + "errors" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" + + coderws "github.com/coder/websocket" + "github.com/stretchr/testify/require" +) + +func TestReadOpenAIWSClientMessage_ControlCloseFrames(t *testing.T) { + tests := []struct { + name string + timeout time.Duration + timeoutStatus coderws.StatusCode + timeoutReason string + cancelCause error + wantStatus coderws.StatusCode + wantReason string + }{ + { + name: "inter-turn idle sends normal close", + timeout: 25 * time.Millisecond, + timeoutStatus: coderws.StatusNormalClosure, + timeoutReason: "websocket idle timeout", + wantStatus: coderws.StatusNormalClosure, + wantReason: "websocket idle timeout", + }, + { + name: "first message timeout sends policy close", + timeout: 25 * time.Millisecond, + timeoutStatus: coderws.StatusPolicyViolation, + timeoutReason: "missing first response.create message", + wantStatus: coderws.StatusPolicyViolation, + wantReason: "missing first response.create message", + }, + { + name: "lease loss sends retry close", + cancelCause: ErrOpenAIWSIngressLeaseLost, + wantStatus: coderws.StatusTryAgainLater, + wantReason: "websocket ingress capacity lease lost; please reconnect", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + controlCtx, cancelControl := context.WithCancelCause(context.Background()) + defer cancelControl(context.Canceled) + serverResult := make(chan error, 1) + readStarted := make(chan struct{}) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + conn, err := coderws.Accept(w, r, nil) + if err != nil { + serverResult <- err + return + } + defer func() { _ = conn.CloseNow() }() + close(readStarted) + _, _, err = ReadOpenAIWSClientMessage( + controlCtx, + conn, + tt.timeout, + tt.timeoutStatus, + tt.timeoutReason, + ) + serverResult <- err + })) + defer server.Close() + + dialCtx, cancelDial := context.WithTimeout(context.Background(), time.Second) + clientConn, _, err := coderws.Dial(dialCtx, "ws"+strings.TrimPrefix(server.URL, "http"), nil) + cancelDial() + require.NoError(t, err) + defer func() { _ = clientConn.CloseNow() }() + <-readStarted + if tt.cancelCause != nil { + cancelControl(tt.cancelCause) + } + + readCtx, cancelRead := context.WithTimeout(context.Background(), time.Second) + _, _, err = clientConn.Read(readCtx) + cancelRead() + var clientClose coderws.CloseError + require.ErrorAs(t, err, &clientClose) + require.Equal(t, tt.wantStatus, clientClose.Code) + require.Equal(t, tt.wantReason, clientClose.Reason) + + select { + case serverErr := <-serverResult: + var closeErr *OpenAIWSClientCloseError + require.ErrorAs(t, serverErr, &closeErr) + require.Equal(t, tt.wantStatus, closeErr.StatusCode()) + require.Equal(t, tt.wantReason, closeErr.Reason()) + case <-time.After(time.Second): + t.Fatal("server read goroutine did not exit after close handshake") + } + }) + } +} + +func TestReadOpenAIWSClientMessage_ParentCancellationStillJoinsRead(t *testing.T) { + controlCtx, cancelControl := context.WithCancelCause(context.Background()) + serverResult := make(chan error, 1) + readStarted := make(chan struct{}) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + conn, err := coderws.Accept(w, r, nil) + if err != nil { + serverResult <- err + return + } + defer func() { _ = conn.CloseNow() }() + close(readStarted) + _, _, err = ReadOpenAIWSClientMessage(controlCtx, conn, 0, 0, "") + serverResult <- err + })) + defer server.Close() + + dialCtx, cancelDial := context.WithTimeout(context.Background(), time.Second) + clientConn, _, err := coderws.Dial(dialCtx, "ws"+strings.TrimPrefix(server.URL, "http"), nil) + cancelDial() + require.NoError(t, err) + defer func() { _ = clientConn.CloseNow() }() + <-readStarted + cancelControl(errors.New("server shutting down")) + readCtx, cancelRead := context.WithTimeout(context.Background(), time.Second) + _, _, err = clientConn.Read(readCtx) + cancelRead() + var clientClose coderws.CloseError + require.ErrorAs(t, err, &clientClose) + require.Equal(t, coderws.StatusGoingAway, clientClose.Code) + require.Equal(t, "websocket request canceled", clientClose.Reason) + + select { + case <-serverResult: + case <-time.After(time.Second): + t.Fatal("server read goroutine leaked after parent cancellation") + } +} diff --git a/backend/internal/service/openai_ws_forwarder_ingress.go b/backend/internal/service/openai_ws_forwarder_ingress.go index 51aace641..0232e7a72 100644 --- a/backend/internal/service/openai_ws_forwarder_ingress.go +++ b/backend/internal/service/openai_ws_forwarder_ingress.go @@ -382,22 +382,18 @@ func (s *OpenAIGatewayService) ProxyResponsesWebSocketFromClient( } readClientMessage := func() ([]byte, error) { - readCtx := ctx idleTimeout := s.openAIWSIngressInterTurnIdleTimeout() - cancelRead := func() {} - if idleTimeout > 0 { - readCtx, cancelRead = context.WithTimeout(ctx, idleTimeout) - } - msgType, payload, readErr := clientConn.Read(readCtx) - cancelRead() + msgType, payload, readErr := ReadOpenAIWSClientMessage( + ctx, + clientConn, + idleTimeout, + coderws.StatusNormalClosure, + "websocket idle timeout", + ) if readErr != nil { - if idleTimeout > 0 && errors.Is(readErr, context.DeadlineExceeded) && ctx.Err() == nil { + var closeErr *OpenAIWSClientCloseError + if errors.As(readErr, &closeErr) && closeErr.StatusCode() == coderws.StatusNormalClosure { logOpenAIWSModeInfo("ingress_ws_inter_turn_idle_timeout account_id=%d timeout_seconds=%d", account.ID, int(idleTimeout.Seconds())) - return nil, NewOpenAIWSClientCloseError( - coderws.StatusNormalClosure, - "websocket idle timeout", - readErr, - ) } return nil, readErr } diff --git a/backend/internal/service/openai_ws_forwarder_ingress_session_test.go b/backend/internal/service/openai_ws_forwarder_ingress_session_test.go index 2578c623d..946a72036 100644 --- a/backend/internal/service/openai_ws_forwarder_ingress_session_test.go +++ b/backend/internal/service/openai_ws_forwarder_ingress_session_test.go @@ -249,6 +249,14 @@ func TestOpenAIGatewayService_ProxyResponsesWebSocketFromClient_IdleTimeoutRelea require.NoError(t, err) require.Equal(t, "response.completed", gjson.GetBytes(event, "type").String()) + closeReadCtx, cancelCloseRead := context.WithTimeout(context.Background(), 3*time.Second) + _, _, err = clientConn.Read(closeReadCtx) + cancelCloseRead() + var clientClose coderws.CloseError + require.ErrorAs(t, err, &clientClose) + require.Equal(t, coderws.StatusNormalClosure, clientClose.Code) + require.Equal(t, "websocket idle timeout", clientClose.Reason) + select { case proxyErr := <-serverErrCh: var closeErr *OpenAIWSClientCloseError diff --git a/backend/internal/service/openai_ws_http_bridge_test.go b/backend/internal/service/openai_ws_http_bridge_test.go index cbc11abcc..808a63c6b 100644 --- a/backend/internal/service/openai_ws_http_bridge_test.go +++ b/backend/internal/service/openai_ws_http_bridge_test.go @@ -775,6 +775,14 @@ func TestOpenAIWSHTTPBridge_IdleTimeoutClosesClientSession(t *testing.T) { require.NoError(t, err) require.Equal(t, "response.completed", gjson.GetBytes(event, "type").String()) + closeReadCtx, cancelCloseRead := context.WithTimeout(context.Background(), 3*time.Second) + _, _, err = clientConn.Read(closeReadCtx) + cancelCloseRead() + var clientClose coderws.CloseError + require.ErrorAs(t, err, &clientClose) + require.Equal(t, coderws.StatusNormalClosure, clientClose.Code) + require.Equal(t, "websocket idle timeout", clientClose.Reason) + select { case proxyErr := <-errCh: var closeErr *OpenAIWSClientCloseError