Skip to content

adds missing cancel() to integration router #3541

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
akshaydeo merged 1 commit into from
May 16, 2026
Merged

adds missing cancel() to integration router #3541

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
29 changes: 12 additions & 17 deletions transports/bifrost-http/integrations/router.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
}
Expand All @@ -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
}
Expand All @@ -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
}
Expand All @@ -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
}
Expand Down Expand Up @@ -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
}
}
Expand All @@ -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
}
Expand Down Expand Up @@ -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"))
Expand All @@ -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"))
Expand All @@ -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"))
Expand All @@ -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"))
Expand All @@ -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"))
Expand All @@ -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
}
Expand All @@ -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
}
Expand All @@ -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)
}
}
Expand Down
Loading