diff --git a/transports/bifrost-http/integrations/router.go b/transports/bifrost-http/integrations/router.go index 476190a149..00ff4274ec 100644 --- a/transports/bifrost-http/integrations/router.go +++ b/transports/bifrost-http/integrations/router.go @@ -667,13 +667,20 @@ func (g *GenericRouter) createHandler(config RouteConfig) fasthttp.RequestHandle // Execute the request through Bifrost bifrostCtx, cancel := lib.ConvertToBifrostContext(ctx, g.handlerStore) + // Centralized cleanup. The streaming branch below transfers ownership via + // streamingOwnsCancel because its producer goroutine outlives this lambda. + streamingOwnsCancel := false + defer func() { + if !streamingOwnsCancel { + cancel() + } + }() // Set integration type to context bifrostCtx.SetValue(schemas.BifrostContextKeyIntegrationType, string(config.Type)) // Async retrieve: check x-bf-async-id header early (before body parsing) if asyncID := string(ctx.Request.Header.Peek(schemas.AsyncHeaderGetID)); asyncID != "" { - defer cancel() g.handleAsyncRetrieve(ctx, config, bifrostCtx) return } @@ -686,7 +693,6 @@ func (g *GenericRouter) createHandler(config RouteConfig) fasthttp.RequestHandle var err error isLargePayload, err = g.largePayloadHook(ctx, bifrostCtx, config.Type) if err != nil { - cancel() g.sendError(ctx, bifrostCtx, config.ErrorConverter, newBifrostError(err, "large payload detection failed")) return } @@ -699,7 +705,6 @@ func (g *GenericRouter) createHandler(config RouteConfig) fasthttp.RequestHandle } else if config.RequestParser != nil { // Use custom parser (e.g., for multipart/form-data) if err := config.RequestParser(ctx, req); err != nil { - cancel() g.sendError(ctx, bifrostCtx, config.ErrorConverter, newBifrostError(err, "failed to parse request")) return } @@ -708,7 +713,6 @@ func (g *GenericRouter) createHandler(config RouteConfig) fasthttp.RequestHandle rawBody = ctx.Request.Body() if len(rawBody) > 0 { if err := sonic.Unmarshal(rawBody, req); err != nil { - cancel() g.sendError(ctx, bifrostCtx, config.ErrorConverter, newBifrostError(err, "Invalid JSON")) return } @@ -753,12 +757,10 @@ func (g *GenericRouter) createHandler(config RouteConfig) fasthttp.RequestHandle if config.ShortCircuit != nil { handled, err := config.ShortCircuit(ctx, bifrostCtx, req) if err != nil { - defer cancel() g.sendError(ctx, bifrostCtx, config.ErrorConverter, newBifrostError(err, "short-circuit handler error: "+err.Error())) return } if handled { - defer cancel() return } } @@ -767,7 +769,6 @@ func (g *GenericRouter) createHandler(config RouteConfig) fasthttp.RequestHandle if config.GetRequestModel != nil { model, err := config.GetRequestModel(ctx, req) if err != nil { - cancel() g.sendError(ctx, bifrostCtx, config.ErrorConverter, newBifrostError(err, "failed to get model from context")) return } @@ -812,7 +813,6 @@ func (g *GenericRouter) createHandler(config RouteConfig) fasthttp.RequestHandle isGenAIBatchCreate := config.Type == RouteConfigTypeGenAI && bifrostCtx.Value(isGeminiBatchCreateRequestContextKey) != nil useBatchPath := config.BatchRequestConverter != nil && (config.RequestConverter == nil || config.Type != RouteConfigTypeGenAI || isGenAIBatchCreate) if useBatchPath { - defer cancel() batchReq, err := config.BatchRequestConverter(bifrostCtx, req) if err != nil { g.sendError(ctx, bifrostCtx, config.ErrorConverter, newBifrostError(err, "failed to convert batch request")) @@ -827,7 +827,6 @@ func (g *GenericRouter) createHandler(config RouteConfig) fasthttp.RequestHandle } // Handle file requests if FileRequestConverter is set if config.FileRequestConverter != nil { - defer cancel() fileReq, err := config.FileRequestConverter(bifrostCtx, req) if err != nil { g.sendError(ctx, bifrostCtx, config.ErrorConverter, newBifrostError(err, "failed to convert file request")) @@ -843,7 +842,6 @@ func (g *GenericRouter) createHandler(config RouteConfig) fasthttp.RequestHandle // Handle container requests if ContainerRequestConverter is set if config.ContainerRequestConverter != nil { - defer cancel() containerReq, err := config.ContainerRequestConverter(bifrostCtx, req) if err != nil { g.sendError(ctx, bifrostCtx, config.ErrorConverter, newBifrostError(err, "failed to convert container request")) @@ -859,7 +857,6 @@ func (g *GenericRouter) createHandler(config RouteConfig) fasthttp.RequestHandle // Handle container file requests if ContainerFileRequestConverter is set if config.ContainerFileRequestConverter != nil { - defer cancel() containerFileReq, err := config.ContainerFileRequestConverter(bifrostCtx, req) if err != nil { g.sendError(ctx, bifrostCtx, config.ErrorConverter, newBifrostError(err, "failed to convert container file request")) @@ -875,7 +872,6 @@ func (g *GenericRouter) createHandler(config RouteConfig) fasthttp.RequestHandle // Handle cached content requests if CachedContentRequestConverter is set if config.CachedContentRequestConverter != nil { - defer cancel() cachedContentReq, err := config.CachedContentRequestConverter(bifrostCtx, req) if err != nil { g.sendError(ctx, bifrostCtx, config.ErrorConverter, newBifrostError(err, "failed to convert cached content request")) @@ -892,12 +888,10 @@ func (g *GenericRouter) createHandler(config RouteConfig) fasthttp.RequestHandle // Convert the integration-specific request to Bifrost format (inference requests) bifrostReq, err := config.RequestConverter(bifrostCtx, req) if err != nil { - defer cancel() g.sendError(ctx, bifrostCtx, config.ErrorConverter, newBifrostError(err, "failed to convert request to Bifrost format")) return } if bifrostReq == nil { - defer cancel() g.sendError(ctx, bifrostCtx, config.ErrorConverter, newBifrostError(nil, "invalid request")) return } @@ -907,14 +901,12 @@ func (g *GenericRouter) createHandler(config RouteConfig) fasthttp.RequestHandle // Extract and parse fallbacks from the request if present if err := g.extractAndParseFallbacks(req, bifrostReq); err != nil { - defer cancel() g.sendError(ctx, bifrostCtx, config.ErrorConverter, newBifrostError(err, "failed to parse fallbacks: "+err.Error())) return } // Async create: check x-bf-async header (needs parsed bifrostReq) if string(ctx.Request.Header.Peek(schemas.AsyncHeaderCreate)) != "" { - defer cancel() g.handleAsyncCreate(ctx, config, req, bifrostReq, bifrostCtx) return } @@ -926,9 +918,12 @@ func (g *GenericRouter) createHandler(config RouteConfig) fasthttp.RequestHandle } if isStreaming { + // Hand cancel ownership to the streaming path; its producer goroutine + // fires cancel on client-disconnect (handleStreaming) and on pre-stream + // errors (handleStreamingRequest). + streamingOwnsCancel = true g.handleStreamingRequest(ctx, config, bifrostReq, bifrostCtx, cancel) } else { - defer cancel() // Ensure cleanup on function exit g.handleNonStreamingRequest(ctx, config, req, bifrostReq, bifrostCtx) } }