fix: cost for image generation or image edit streaming

[sammaji]

Summary

Eliminates the intermediate collectedUsage variable in image generation and image edit streaming handlers, instead using response.Usage directly on the completed event. Also preserves any OutputTokensDetails already present in the response rather than unconditionally overwriting it, and hides NumInputImages from JSON serialization.

Changes

• Removed the collectedUsage accumulator variable from both HandleOpenAIImageGenerationStreaming and HandleOpenAIImageEditStreamRequest. Usage data is now read directly from response.Usage when the completed event is received, avoiding an unnecessary copy.
• Updated the NImages backfill logic to only initialize OutputTokensDetails if it is nil, and only set NImages if it is not already populated — preserving any details the provider may have already returned.
• Changed NumInputImages in ImageUsage to use json:"-" so it is no longer included in serialized responses (it is an internal Bifrost field).
• Changed ExtraFields in BifrostImageGenerationResponse and BifrostImageGenerationStreamResponse to always serialize (json:"extra_fields" instead of json:"extra_fields,omitempty").
• Removed the unused getModelFromRequest helper function from schemas/images.go.

Closes #4777

Type of change

• Refactor
• Bug fix

Affected areas

• Core (Go)
• Providers/Integrations

How to test

go test ./...

Trigger a streaming image generation and image edit request and verify that:

• Usage fields are correctly populated on the completed chunk.
• NImages reflects the number of images generated.
• num_input_images no longer appears in the JSON response.
• extra_fields is always present in the response body.

Breaking changes

• No

Related issues

Security considerations

None.

Checklist

• I read docs/contributing/README.md and followed the guidelines
• I added/updated tests where appropriate
• I updated documentation where needed
• I verified builds succeed (Go and UI)
• I verified the CI pipeline passes locally if applicable

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro Plus

Run ID: 1510ee62-b537-4ace-8ec2-2eee205e16a5

📥 Commits

Reviewing files that changed from the base of the PR and between 28a3c79 and a5285a5.

📒 Files selected for processing (2)

• core/providers/openai/openai.go
• core/schemas/images.go

🚧 Files skipped from review as they are similar to previous changes (2)

• core/schemas/images.go
• core/providers/openai/openai.go

---

📝 Walkthrough

Summary by CodeRabbit

• Bug Fixes
  • Improved streaming image generation/edit usage reporting so completed chunks reliably return correct token and image counts in final results.
  • Updated streaming completion detection to treat both generation and edit completion events as finished.
• Changes
  • Ensured extra_fields is always included in image generation responses.
  • Stopped returning the input image count in image usage output.
  • Improved automatic model backfilling when a model is not provided in image generation requests.

Walkthrough

Updates image streaming completion handling to use streamed usage data directly, backfill missing image counts, and adjust image response schema serialization and backfill behavior.

Changes

Image Streaming Usage Fix

Layer / File(s)	Summary
Schema JSON and backfill changes core/schemas/images.go	ExtraFields is always emitted on both image response structs, NumInputImages is internal-only, and model backfilling is handled inline in BackfillParams.
Streaming usage completion handling core/providers/openai/openai.go	Removes collectedUsage accumulation, treats both completed event types as completion, ensures OutputTokensDetails exists, backfills NImages when missing, and assigns chunk.Usage from response.Usage.

Estimated code review effort: 2 (Simple) | ~10 minutes

Possibly related PRs

• maximhq/bifrost#3495: Modifies the same OpenAI image streaming handlers and related control flow in core/providers/openai/openai.go.

Suggested reviewers: akshaydeo

🚥 Pre-merge checks | ✅ 3 | ❌ 2

❌ Failed checks (2 warnings)

Check name	Status	Explanation	Resolution
Out of Scope Changes check	⚠️ Warning	Schema serialization changes and helper removal go beyond #4777's streaming cost fix and appear unrelated to the linked issue.	Split unrelated schema cleanup and JSON-tag changes into a separate PR, or justify them as required for the #4777 fix.
Docstring Coverage	⚠️ Warning	Docstring coverage is 33.33% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (3 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly matches the main change: fixing image generation and image edit streaming cost handling.
Description check	✅ Passed	The description follows the template and covers summary, changes, testing, type, areas, breaking changes, related issues, and checklist.
Linked Issues check	✅ Passed	The PR addresses #4777 by preserving output_tokens_details and populating NImages on completed streamed image responses.

✨ Finishing Touches

📝 Generate docstrings

• Create stacked PR
• Commit on current branch

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch 06-30-fix_cost_for_image_generation___edit_streaming

Warning

There were issues while running some tools. Please review the errors and either fix the tool's configuration or disable the tool if it's a critical failure.

🔧 golangci-lint (2.12.2)

level=error msg="[linters_context] typechecking error: pattern ./...: directory prefix . does not contain main module or its selected dependencies"

---

_{Comment @coderabbitai help to get the list of available commands.}

[sammaji]

• fix: cost for image generation or image edit streaming #4802 👈 (View in Graphite)
• dev

This stack of pull requests is managed by Graphite. Learn more about stacking.

[greptile-apps]

Confidence Score: 5/5

Safe to merge — the change correctly targets a narrowly scoped streaming usage bug with no mutations to shared state.

The refactor removes an intermediate copy that was silently dropping InputTokensDetails, and replaces it with a direct pointer assignment that is safe because the cost calculation path already guards it with DeepCopy(). The NImages backfill guard, json:"-" on NumInputImages, and extra_fields omitempty removal are all consistent with the rest of the codebase and do not introduce regressions.

No files require special attention.

Important Files Changed

Filename	Overview
core/providers/openai/openai.go	Removes collectedUsage accumulator from both image streaming handlers; usage is now read directly from response.Usage on the completed event, fixing silent InputTokensDetails loss. NImages backfill is now guarded by == 0 to preserve provider-supplied values. The aliasing (chunk.Usage = response.Usage) is safe because cost calculation calls DeepCopy() and response is a per-iteration local.
core/schemas/images.go	Three schema changes: NumInputImages tagged json:"-" (was omitempty); ExtraFields on both image response structs loses omitempty to match every other Bifrost response type; dead helper getModelFromRequest removed.

Sequence Diagram

%%{init: {'theme': 'neutral'}}%%
sequenceDiagram
    participant OAI as OpenAI SSE Stream
    participant Handler as HandleOpenAIImage*Streaming
    participant Chunk as BifrostImageGenerationStreamResponse
    participant PostHook as Post-hook / Cost Calc

    OAI->>Handler: partial_image chunk
    Handler->>Handler: update maxImageIndex
    Handler->>Chunk: build chunk (no usage)
    Handler->>PostHook: ProcessAndSendResponse(chunk)

    OAI->>Handler: completed chunk (with Usage)
    Handler->>Handler: update maxImageIndex
    Handler->>Handler: "if OutputTokensDetails==nil → allocate"
    Handler->>Handler: "if NImages==0 → set maxImageIndex+1"
    Handler->>Chunk: "chunk.Usage = response.Usage (pointer)"
    Handler->>Chunk: BackfillParams → sets NumInputImages
    Handler->>PostHook: ProcessAndSendResponse(chunk)
    PostHook->>PostHook: chunk.Usage.DeepCopy() → independent copy
    PostHook->>PostHook: compute cost (NImages, NumInputImages, InputTokensDetails preserved)

Loading

%%{init: {'theme': 'base', 'themeVariables': {"darkMode": true, "background": "#0d1117", "primaryColor": "#21262d", "primaryTextColor": "#e6edf3", "primaryBorderColor": "#8b949e", "lineColor": "#8b949e", "textColor": "#e6edf3", "edgeLabelBackground": "#161b22", "actorBkg": "#21262d", "actorBorder": "#8b949e", "actorTextColor": "#e6edf3", "actorLineColor": "#8b949e", "signalColor": "#8b949e", "signalTextColor": "#e6edf3", "noteBkgColor": "#373320", "noteBorderColor": "#d4a72c", "noteTextColor": "#f0e6c0", "labelBoxBkgColor": "#21262d", "labelBoxBorderColor": "#8b949e", "labelTextColor": "#e6edf3", "loopTextColor": "#e6edf3", "activationBkgColor": "#30363d", "activationBorderColor": "#8b949e"}}}%%
sequenceDiagram
    participant OAI as OpenAI SSE Stream
    participant Handler as HandleOpenAIImage*Streaming
    participant Chunk as BifrostImageGenerationStreamResponse
    participant PostHook as Post-hook / Cost Calc

    OAI->>Handler: partial_image chunk
    Handler->>Handler: update maxImageIndex
    Handler->>Chunk: build chunk (no usage)
    Handler->>PostHook: ProcessAndSendResponse(chunk)

    OAI->>Handler: completed chunk (with Usage)
    Handler->>Handler: update maxImageIndex
    Handler->>Handler: "if OutputTokensDetails==nil → allocate"
    Handler->>Handler: "if NImages==0 → set maxImageIndex+1"
    Handler->>Chunk: "chunk.Usage = response.Usage (pointer)"
    Handler->>Chunk: BackfillParams → sets NumInputImages
    Handler->>PostHook: ProcessAndSendResponse(chunk)
    PostHook->>PostHook: chunk.Usage.DeepCopy() → independent copy
    PostHook->>PostHook: compute cost (NImages, NumInputImages, InputTokensDetails preserved)

Loading

_{Reviews (5): Last reviewed commit: "fix: cost for image generation and image..." | Re-trigger Greptile}

The merge-base changed after approval.

[CLAassistant]

All committers have signed the CLA.

[akshaydeo]

Merge activity

• Jul 3, 2:17 PM UTC: A user started a stack merge that includes this pull request via Graphite.
• Jul 3, 2:17 PM UTC: @akshaydeo merged this pull request with Graphite.