Skip to content

refactor(virtualmodels): unify batch-rewrite paths and drop dead rewrite mode#424

Merged
SantiagoDePolonia merged 1 commit into
mainfrom
refactor/virtualmodels-batch-dry
Jun 23, 2026
Merged

refactor(virtualmodels): unify batch-rewrite paths and drop dead rewrite mode#424
SantiagoDePolonia merged 1 commit into
mainfrom
refactor/virtualmodels-batch-dry

Conversation

@SantiagoDePolonia

@SantiagoDePolonia SantiagoDePolonia commented Jun 23, 2026

Copy link
Copy Markdown
Contributor

Summary

DRY follow-up to #423 (review item #3). Now that #423 has merged, this PR is rebased onto main and keeps only the batch-rewrite cleanup.

  • Unifies provider-wrapper and server-side batch rewriting behind rewriteBatchSource / rewriteBatchItem.
  • BatchPreparer passes ValidateModelAccess; the provider wrapper passes nil, preserving existing wrapper behavior.
  • Removes the now-dead requestRewriteMode / rewriteForUpstream machinery from translated request rewriting.
  • Adds behavior coverage for provider clearing, access denial, nil validation, cross-provider batch rejection, and unsupported batch items.

Verification

  • go test ./internal/virtualmodels
  • go test ./...
  • Pre-commit hook: make test-race, make lint, go mod tidy

🤖 Generated with Claude Code

Summary by CodeRabbit

  • Refactor
    • Streamlined batch request rewriting for redirected models, ensuring the resolved provider is consistently forwarded for upstream processing.
    • Centralized per-item access checks and standardized rewritten batch-item handling, including encoding errors.
    • Simplified request rewriting for chat, responses, and embeddings to always preserve the resolved provider.
  • Bug Fixes
    • Malformed or unsupported batch items are rejected.
    • Prevents unauthorized rewrites and rejects cross-provider batches when the resolved target provider differs.
  • Tests
    • Added coverage for batch item rewriting, access-control enforcement, and invalid/cross-provider scenarios.

@coderabbitai

coderabbitai Bot commented Jun 23, 2026

Copy link
Copy Markdown
Contributor

Review Change Stack

Note

Reviews paused

It looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the reviews.auto_review.auto_pause_after_reviewed_commits setting.

Use the following commands to manage reviews:

  • @coderabbitai resume to resume automatic reviews.
  • @coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

  • ▶️ Resume reviews
  • 🔍 Trigger review
📝 Walkthrough

Walkthrough

Two files in internal/virtualmodels are refactored. provider.go removes the requestRewriteMode type and constants, drops the github.com/goccy/go-json import, and simplifies rewriteChatRequest, rewriteResponsesRequest, and rewriteEmbeddingRequest to always forward the resolved provider. batch_preparer.go extracts inline batch rewriting into reusable helpers: rewriteBatchSource, rewriteBatchItem, validateAccess, batchFileTransport, and marshalBatchItem. All batch call sites are updated with the new signature, and batch preparation is validated by a comprehensive test suite.

Changes

Rewrite Helper Simplification and Batch Extraction

Layer / File(s) Summary
Remove requestRewriteMode and replace rewrite helpers
internal/virtualmodels/provider.go
Deletes requestRewriteMode type and constants, removes the github.com/goccy/go-json import, and replaces rewriteChatRequest, rewriteResponsesRequest, and rewriteEmbeddingRequest with new versions accepting only (ctx, service, checker, req) that resolve redirects with an empty expected-provider-type and always set forward.Provider from the resolved selector.
Update request rewrite call sites
internal/virtualmodels/provider.go
All five request-rewriting call sites in ChatCompletion, StreamChatCompletion, Responses, StreamResponses, and Embeddings are updated to use the reduced helper signature without rewrite-mode parameters.
Extract batch preparation helpers
internal/virtualmodels/batch_preparer.go
PrepareBatchRequest now delegates to rewriteBatchSource, which uses per-item rewriteBatchItem function. New helpers centralize concerns: validateAccess enforces model access when configured; batchFileTransport selects native file providers; rewriteBatchItem resolves redirects, validates access, and rewrites items; marshalBatchItem wraps JSON encoding with proper error handling.
Update batch method calls
internal/virtualmodels/provider.go
CreateBatch, CreateBatchWithHints, and PrepareBatchRequest each update their rewriteBatchSource call signatures to pass file transport and a trailing nil argument for the optional validate callback.
Test batch item rewriting and validation
internal/virtualmodels/batch_preparer_test.go
Adds comprehensive test suite verifying that rewriting resolves redirected models and clears providers, that validation hooks reject unauthorized requests, that unsupported batch item types are rejected, that cross-provider batches fail, and that validateAccess enforces access policy for disabled virtual models while allowing nil-service preparers.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~25 minutes

Poem

🐇 Hop, hop, the mode is gone,
No more constants to rely upon!
Helpers extracted, neat and clean,
The batch now flows through functions lean.
Marshal errors wrapped with care,
Tests prove the flow is standing fair! ✨

🚥 Pre-merge checks | ✅ 5
✅ Passed checks (5 passed)
Check name Status Explanation
Description Check ✅ Passed Check skipped - CodeRabbit’s high-level summary is enabled.
Title check ✅ Passed The title directly and concisely describes the main refactoring: unifying batch-rewrite paths and removing dead rewrite mode code.
Docstring Coverage ✅ Passed Docstring coverage is 81.25% which is sufficient. The required threshold is 80.00%.
Linked Issues check ✅ Passed Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check ✅ Passed Check skipped because no linked issues were found for this pull request.

✏️ Tip: You can configure your own custom pre-merge checks in the settings.

✨ Finishing Touches
🧪 Generate unit tests (beta)
  • Create PR with unit tests
  • Commit unit tests in branch refactor/virtualmodels-batch-dry

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

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

@greptile-apps

greptile-apps Bot commented Jun 23, 2026

Copy link
Copy Markdown

Confidence Score: 5/5

This looks safe to merge.

  • No blocking issues found in the changed code.

Reviews (2): Last reviewed commit: "test(virtualmodels): cover the unified b..." | Re-trigger Greptile

Comment thread internal/virtualmodels/batch_preparer.go
@SantiagoDePolonia

Copy link
Copy Markdown
Contributor Author

Thanks @greptile-apps — added behavior tests for the unified batch path in internal/virtualmodels/batch_preparer_test.go (2a568b2), covering all four you called out plus a couple more:

  • Server-side validation on the resolved selectorrewriteBatchItem with a denying validate hook surfaces the error; BatchPreparer.validateAccess denies a disabled model.
  • Provider-wrapper nil validationrewriteBatchItem(..., nil) rewrites without an access check; a nil-service BatchPreparer never blocks.
  • Provider cleared in the rewritten body — asserts the resolved redirect is written as the model with provider cleared for upstream.
  • Single-provider-per-batch — a resolved target whose provider differs from the batch provider is rejected.
  • Unsupported/malformed items — rejected rather than passed through.

I also brought this branch up to date with its base (#423), which resolved the overlap with the scoped-redirect fix landed there — rewriteBatchItem now resolves via the user-path-aware resolveRedirectRoutableSelector, so batch requests honor scoped aliases too.

go build/vet/golangci-lint (0 issues) / go test ./... (57 pkg) all green.

@greptile-apps

greptile-apps Bot commented Jun 23, 2026

Copy link
Copy Markdown

Tests look great — solid coverage of all the key paths:

  • Nil-validate + provider cleared (TestRewriteBatchItem_RewritesAndClearsProvider): correctly verifies gpt-4o with empty provider on the rewritten body.
  • Access denial (TestRewriteBatchItem_ValidateRejectsUnauthorized): uses errors.Is so wrapped errors are handled correctly.
  • Cross-provider rejection (TestRewriteBatchItem_RejectsCrossProviderBatch): the "anthropic" providerType against an openai-resolving redirect is the right minimal setup for that invariant.
  • Unsupported item (TestRewriteBatchItem_UnsupportedItem): passing a raw string as Request hits the default branch cleanly.
  • validateAccess nil-service parity (TestBatchPreparerValidateAccess): the two sub-cases (service with disabled model, zero-value BatchPreparer) map exactly to the two branches in the method.

One small note: TestBatchPreparerValidateAccess upserts openai/gpt-4o with Enabled: false but then calls validateAccess with a raw core.ModelSelector{Provider: "openai", Model: "gpt-4o"} — worth confirming ValidateModelAccess matches on that selector form directly (rather than requiring it to be resolved through the redirect table first). If newTestService already seeds the catalog such that openai/gpt-4o is a known model and the policy lookup is by resolved selector, this is fine — just something to double-check if the test starts flaking under catalog changes.

Otherwise the implementation is clean and the tests are thorough. LGTM.

@SantiagoDePolonia SantiagoDePolonia force-pushed the refactor/virtualmodels-batch-dry branch from 2a568b2 to e6b66d7 Compare June 23, 2026 14:55
@SantiagoDePolonia SantiagoDePolonia changed the base branch from feat/unifying-ui-for-virtual-models to main June 23, 2026 14:55

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Actionable comments posted: 1

🤖 Prompt for all review comments with 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.

Inline comments:
In `@internal/virtualmodels/batch_preparer.go`:
- Around line 88-119: The function `requestedSelectorForDecodedRequest`
duplicates the logic already provided by
`core.DecodedBatchItemRequest.RequestedModelSelector()` method, which properly
handles selector extraction for all three supported request types using the
`semanticSelectorCarrier` interface. Replace the call to
`requestedSelectorForDecodedRequest(decoded.Request)` with a direct call to
`decoded.RequestedModelSelector()` to eliminate the redundant type switch and
error reconstruction at line 90. After making this change, remove the entire
`requestedSelectorForDecodedRequest` function since it will no longer be needed.
🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

  • Push a commit to this branch (recommended)
  • Create a new PR with the fixes

ℹ️ Review info
⚙️ Run configuration

Configuration used: Organization UI

Review profile: ASSERTIVE

Plan: Pro

Run ID: e1f360a5-96da-40cb-8672-d77495c20977

📥 Commits

Reviewing files that changed from the base of the PR and between 2a568b2 and e6b66d7.

📒 Files selected for processing (3)
  • internal/virtualmodels/batch_preparer.go
  • internal/virtualmodels/batch_preparer_test.go
  • internal/virtualmodels/provider.go

Comment thread internal/virtualmodels/batch_preparer.go Outdated
@SantiagoDePolonia SantiagoDePolonia force-pushed the refactor/virtualmodels-batch-dry branch from e6b66d7 to b9d304a Compare June 23, 2026 15:45
@codecov-commenter

codecov-commenter commented Jun 23, 2026

Copy link
Copy Markdown

⚠️ Please install the 'codecov app svg image' to ensure uploads and comments are reliably processed by Codecov.

Codecov Report

❌ Patch coverage is 57.44681% with 20 lines in your changes missing coverage. Please review.

Files with missing lines Patch % Lines
internal/virtualmodels/batch_preparer.go 60.00% 11 Missing and 1 partial ⚠️
internal/virtualmodels/provider.go 52.94% 8 Missing ⚠️

📢 Thoughts on this report? Let us know!

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Actionable comments posted: 4

🤖 Prompt for all review comments with 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.

Inline comments:
In `@internal/virtualmodels/batch_preparer_test.go`:
- Around line 87-102: The TestBatchPreparerValidateAccess test function
currently only validates the failure cases (disabled model and nil service) but
lacks coverage for the success path. Add a test case within this function that
verifies validateAccess succeeds when called with a non-nil BatchPreparer
service and an enabled model. This should mirror the existing test structure by
first upserting an enabled VirtualModel with the same provider and model
selector, then calling validateAccess on a properly initialized BatchPreparer
instance and asserting that no error is returned, thereby completing the
validation behavior matrix.
- Around line 78-83: The TestRewriteBatchItem_RejectsCrossProviderBatch test
currently only verifies that an error is returned but does not check the error
message content. Modify the test to assert that the returned error message
contains relevant keywords like "single-provider" or "cross-provider" to verify
the correct rejection reason is being enforced. This can be done by checking
that the error message string contains the expected substring after confirming
err is not nil, making the test more self-documenting and better at catching
regressions where the wrong error type is returned.
- Around line 96-98: The test for NewBatchPreparer's validateAccess method when
called with a disabled model only verifies that an error is returned, but does
not verify the error type and code. To ensure compliance with the
ValidateModelAccess contract, enhance the test to assert that the returned error
is specifically an InvalidRequestError with a "model_access_denied" code. This
can be done by type-asserting the error and checking its code field to match the
expected denial response.
- Around line 68-74: In the TestRewriteBatchItem_UnsupportedItem test function,
the current error check only verifies that an error is returned when
rewriteBatchItem is called with an unsupported endpoint. To improve regression
detection, extend the test to also verify the error type matches the expected
invalid request error. Replace the simple nil check with a type assertion or
error comparison that confirms the returned error is the correct invalid request
error type expected for unsupported batch items.
🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

  • Push a commit to this branch (recommended)
  • Create a new PR with the fixes

ℹ️ Review info
⚙️ Run configuration

Configuration used: Organization UI

Review profile: ASSERTIVE

Plan: Pro

Run ID: 8effaae7-f4ee-41a9-ab4e-0a72813d55e7

📥 Commits

Reviewing files that changed from the base of the PR and between e6b66d7 and b9d304a.

📒 Files selected for processing (3)
  • internal/virtualmodels/batch_preparer.go
  • internal/virtualmodels/batch_preparer_test.go
  • internal/virtualmodels/provider.go

Comment thread internal/virtualmodels/batch_preparer_test.go
Comment thread internal/virtualmodels/batch_preparer_test.go
Comment thread internal/virtualmodels/batch_preparer_test.go
Comment thread internal/virtualmodels/batch_preparer_test.go Outdated
@SantiagoDePolonia SantiagoDePolonia force-pushed the refactor/virtualmodels-batch-dry branch 2 times, most recently from d45d072 to ad7efab Compare June 23, 2026 17:15

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Actionable comments posted: 2

🤖 Prompt for all review comments with 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.

Inline comments:
In `@internal/virtualmodels/batch_preparer_test.go`:
- Around line 83-90: The test TestRewriteBatchItem_ValidateRejectsUnauthorized
only verifies that the validator hook error is surfaced, but does not verify
that the validator receives the correct resolved selector. Strengthen the test
by capturing the core.ModelSelector argument passed to the validator function
and asserting it matches the expected resolved target (openai/gpt-4o) rather
than the requested alias (fast). This ensures validation is operating on the
correctly translated model reference, not just that errors propagate.

In `@internal/virtualmodels/provider.go`:
- Around line 528-572: Add focused unit tests for the three rewrite functions:
rewriteChatRequest, rewriteResponsesRequest, and rewriteEmbeddingRequest. Each
test should cover three aspects: successful redirect resolution where the
selector model and provider are correctly rewritten into the returned forward
request, error propagation when resolveRedirectRoutableSelector fails, and nil
request handling. These tests will ensure the provider preservation behavior for
routing and prevent regression against the batch-clearing behavior documented in
the comments.
🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

  • Push a commit to this branch (recommended)
  • Create a new PR with the fixes

ℹ️ Review info
⚙️ Run configuration

Configuration used: Organization UI

Review profile: ASSERTIVE

Plan: Pro

Run ID: 63950b4c-379c-4603-951b-e0d22bd5af0e

📥 Commits

Reviewing files that changed from the base of the PR and between b9d304a and d45d072.

📒 Files selected for processing (3)
  • internal/virtualmodels/batch_preparer.go
  • internal/virtualmodels/batch_preparer_test.go
  • internal/virtualmodels/provider.go

Comment thread internal/virtualmodels/batch_preparer_test.go
Comment thread internal/virtualmodels/provider.go
@SantiagoDePolonia SantiagoDePolonia force-pushed the refactor/virtualmodels-batch-dry branch from ad7efab to 0f472f0 Compare June 23, 2026 17:41
@SantiagoDePolonia SantiagoDePolonia merged commit 7504016 into main Jun 23, 2026
18 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants