Skip to content

fix(passthrough): track usage for non-streaming responses and extract… #332

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

Open
vfeitoza wants to merge 1 commit into
base: main
Choose a base branch
from
+48 −2
Open

fix(passthrough): track usage for non-streaming responses and extract… #332

Show file tree
Hide file tree
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
45 changes: 44 additions & 1 deletion internal/server/passthrough_support.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -311,12 +311,55 @@ func (s *passthroughService) proxyPassthroughResponse(c *echo.Context, providerT
return nil
}

body, err := io.ReadAll(resp.Body)
if err != nil {
return handleError(c, core.NewProviderError(providerType, http.StatusBadGateway, "failed to read provider passthrough response body", err))
}
Comment on lines +314 to +317
Copy link
Copy Markdown
Contributor

@coderabbitai coderabbitai Bot May 15, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major | 🏗️ Heavy lift

Avoid unbounded buffering of non-SSE passthrough responses.

Line 314 reads the full upstream body into memory with io.ReadAll(resp.Body) and no size guard. Large passthrough payloads can cause high memory pressure or OOM under concurrency.

🤖 Prompt for AI Agents 
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@internal/server/passthrough_support.go` around lines 314 - 317, The code
currently does io.ReadAll(resp.Body) which can unboundedly buffer large
passthrough responses; replace this with a bounded-read or streaming approach:
for non-SSE passthrough responses use an io.LimitedReader (e.g.,
io.LimitReader(resp.Body, maxPassthroughBytes)) or stream directly to the client
with io.Copy (copying from resp.Body to the framework response writer) and
return a ProviderError if the limit is exceeded. Update the code around
resp.Body read (the block that calls io.ReadAll and returns handleError(...)) to
enforce a configurable maxPassthroughBytes constant and avoid loading the entire
body into memory. Ensure you still close resp.Body and preserve existing error
handling via handleError/providerType.


workflow := core.GetWorkflow(c.Request().Context())
if s.usageLogger != nil && s.usageLogger.Config().Enabled && (workflow == nil || workflow.UsageEnabled()) {
model := ""
if info != nil {
model = strings.TrimSpace(info.Model)
}
model = resolvedModelFromWorkflow(workflow, model)
requestID := requestIDFromContextOrHeader(c.Request())
usagePath := strings.TrimSpace(c.Request().URL.Path)
s.logPassthroughNonStreamUsage(body, model, providerType, providerName, requestID, usagePath, c.Request().Context())
}
Comment on lines +327 to +329
Copy link
Copy Markdown
Contributor

@coderabbitai coderabbitai Bot May 15, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major | ⚡ Quick win

Fix non-stream audit path derivation to use provider endpoint, not request URL path.

At Line 328, the caller passes request.URL.Path into logPassthroughNonStreamUsage as endpoint. At Line 346, that value is used as both requestPath and endpoint in passthroughStreamAuditPath(...), which prevents canonical mapping (e.g., /chat/completions → /v1/chat/completions) and mislabels usage entries.

Suggested minimal fix 
-        usagePath := strings.TrimSpace(c.Request().URL.Path)
-        s.logPassthroughNonStreamUsage(body, model, providerType, providerName, requestID, usagePath, c.Request().Context())
+        requestPath := strings.TrimSpace(c.Request().URL.Path)
+        s.logPassthroughNonStreamUsage(body, model, providerType, providerName, requestID, requestPath, endpoint, c.Request().Context())
@@
-func (s *passthroughService) logPassthroughNonStreamUsage(body []byte, model, providerType, providerName, requestID, endpoint string, ctx context.Context) {
+func (s *passthroughService) logPassthroughNonStreamUsage(body []byte, model, providerType, providerName, requestID, requestPath, providerEndpoint string, ctx context.Context) {
@@
-    auditPath := passthroughStreamAuditPath(endpoint, providerType, endpoint)
+    auditPath := passthroughStreamAuditPath(requestPath, providerType, providerEndpoint)

Also applies to: 341-347

🤖 Prompt for AI Agents 
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@internal/server/passthrough_support.go` around lines 327 - 329, The code
currently passes the incoming request URL path (usagePath) into
logPassthroughNonStreamUsage which later uses it as both requestPath and
endpoint in passthroughStreamAuditPath, preventing canonical endpoint mapping;
change logPassthroughNonStreamUsage calls (and the values passed into
passthroughStreamAuditPath) to pass two distinct values: the request path
(strings.TrimSpace(c.Request().URL.Path)) as requestPath and the provider's
configured endpoint (the provider endpoint obtained from the provider config
used for the outbound call) as endpoint; update calls around
logPassthroughNonStreamUsage and passthroughStreamAuditPath (also in the 341-347
region) to use providerEndpoint for endpoint so canonical mapping (e.g.,
/v1/chat/completions) is preserved.


c.Response().WriteHeader(resp.StatusCode)
if _, err := io.Copy(c.Response(), resp.Body); err != nil {
if _, err := c.Response().Write(body); err != nil {
return err
}
if f, ok := c.Response().(http.Flusher); ok {
f.Flush()
}
return nil
}

func (s *passthroughService) logPassthroughNonStreamUsage(body []byte, model, providerType, providerName, requestID, endpoint string, ctx context.Context) {
if len(body) == 0 {
return
}

auditPath := passthroughStreamAuditPath(endpoint, providerType, endpoint)
Copy link
Copy Markdown

@greptile-apps greptile-apps Bot May 15, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 The endpoint received by logPassthroughNonStreamUsage is c.Request().URL.Path (e.g., /p/openai/v1/chat/completions), not the passthrough endpoint segment (e.g., chat/completions). Passing it as both requestPath and endpoint means passthroughStreamAuditPath's switch cases (/chat/completions, /responses, /messages) never match, so auditPath always equals the full URL path. The streaming path avoids this by calling passthroughAuditPath(c, providerType, endpoint, info) with the real endpoint segment, which correctly returns /v1/chat/completions for OpenAI chat. The non-streaming path should derive the audit path the same way before calling the helper, rather than relying on the URL path for normalisation.

Suggested change
auditPath := passthroughStreamAuditPath(endpoint, providerType, endpoint)
auditPath := passthroughStreamAuditPath(endpoint, providerType, endpoint) // TODO: pass the actual passthrough endpoint segment, not the full URL path

var pricingArgs []*core.ModelPricing
if s.pricingResolver != nil {
pricingProvider := strings.TrimSpace(providerName)
if pricingProvider == "" {
pricingProvider = strings.TrimSpace(providerType)
}
if p := s.pricingResolver.ResolvePricing(model, pricingProvider); p != nil {
pricingArgs = append(pricingArgs, p)
}
}

entry := usage.ExtractFromCachedResponseBody(body, requestID, model, providerType, auditPath, "", pricingArgs...)
if entry == nil {
return
}
entry.ProviderName = strings.TrimSpace(providerName)
entry.UserPath = core.UserPathFromContext(ctx)
s.usageLogger.Write(entry)
}
5 changes: 4 additions & 1 deletion internal/server/request_selector_peek.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,10 @@ func seedRequestBodySelectorHints(req *http.Request, bodyMode core.BodyMode, env
}

hints := peekRequestBodySelectorHints(req, requestSelectorPeekLimit)
if !hints.parsed || !hints.complete {
if !hints.parsed && !hints.complete {
if bodyMode == core.BodyModeOpaque && hints.model != "" {
core.ApplyBodySelectorHints(env, hints.model, hints.provider, hints.stream)
}
return
}
Comment on lines +29 to 34
Copy link
Copy Markdown

@greptile-apps greptile-apps Bot May 15, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 || → && broadens the fall-through path beyond the described fix

The change from !hints.parsed || !hints.complete to !hints.parsed && !hints.complete also silently changes behaviour for a second state: when the decoder finds both model and provider and exits early, it returns parsed=true, complete=false. The old guard returned without applying hints in that case (a pre-existing bug); the new guard falls through to core.ApplyBodySelectorHints, which is correct. This second fix is beneficial, but it is not mentioned in the PR description — worth calling out explicitly so the change is clearly intentional.

core.ApplyBodySelectorHints(env, hints.model, hints.provider, hints.stream)
Expand Down