feat(proxy): move chat completion to new provider integration#37
Merged
feat(proxy): move chat completion to new provider integration#37
Conversation
📝 WalkthroughWalkthroughSwitched Usage total-token derivation to centralized helpers and updated callers to call Changes
Sequence Diagram(s)sequenceDiagram
participant Client
participant Handler as Chat Handler
participant Gateway
participant Provider as Provider Instance
participant RateLimit as Rate Limiter
participant Observability
Client->>Handler: POST /chat/completions
Handler->>Gateway: resolve provider instance (model)
Gateway-->>Handler: ProviderInstance
Handler->>Provider: chat(request, provider_instance)
alt Streaming
Provider-->>Handler: ChatResponseStream
Handler->>Handler: spawn_stream_usage_observer()
Handler-->>Client: SSE stream chunks
Handler->>RateLimit: post_check_streaming(usage)
Handler->>Observability: record_streaming_usage(usage)
else Non-Streaming
Provider-->>Handler: Complete response (includes Usage)
Handler->>RateLimit: post_check(usage)
Handler->>Observability: record_usage(usage)
Handler-->>Client: JSON response
end
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Possibly related PRs
Important Pre-merge checks failedPlease resolve all errors before merging. Addressing warnings is optional. ❌ Failed checks (1 error, 1 warning)
✅ Passed checks (2 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (4)
src/gateway/traits/provider.rs (1)
324-332: Test covers the primary use case.Consider adding edge case tests for robustness (optional):
build_url("https://example.com", "ignored")→ no overlapbuild_url("https://example.com/v1/", "ignored")→ trailing slash handlingbuild_url("not-a-valid-url", "ignored")→ fallback path🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/gateway/traits/provider.rs` around lines 324 - 332, Add additional unit tests for build_url on DummyProvider to cover edge cases: assert build_url("https://example.com", "ignored") returns "https://example.com/chat/completions" (no overlapping version), assert build_url("https://example.com/v1/", "ignored") handles trailing slash and returns "https://example.com/v1/chat/completions", and assert build_url("not-a-valid-url", "ignored") exercises the fallback path behavior; add these alongside the existing build_url_avoids_duplicate_version_prefixes test and ensure the tests reference DummyProvider::build_url to locate the implementation.src/proxy/hooks/observability/mod.rs (1)
61-90: Trace the public observability hooks.
record_start_time,record_usage,record_first_token_latency, andrecord_streaming_usageare public entrypoints undersrc/, but none has#[fastrace::trace].As per coding guidelines,
src/**/*.rs: Use #[fastrace::trace] decorator for distributed tracing on public functions.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/proxy/hooks/observability/mod.rs` around lines 61 - 90, Add the #[fastrace::trace] attribute to each public observability hook to enable distributed tracing: place #[fastrace::trace] directly above the async fn declarations for record_start_time, record_usage, record_first_token_latency, and record_streaming_usage; ensure the fastrace macro is in scope (import if necessary) so the attribute compiles cleanly.src/proxy/provider.rs (1)
14-21: Document and trace the new public provider helpers.
create_legacy_providerandcreate_provider_instanceare exported fromsrc/, but they are missing///docs and#[fastrace::trace].As per coding guidelines,
src/**/*.rs: Use #[fastrace::trace] decorator for distributed tracing on public functions, and**/*.rs: Use /// for doc comments on public items in Rust.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/proxy/provider.rs` around lines 14 - 21, Add missing documentation comments and fastrace decorators to the public helper functions: annotate create_legacy_provider and create_provider_instance with a brief triple-slash doc comment describing purpose, parameters (e.g., model, gateway) and return value, and add the #[fastrace::trace] attribute above each function to enable distributed tracing; ensure the docs mention that create_legacy_provider returns a Box<dyn Provider> created from model.provider_config and that create_provider_instance returns Result<ProviderInstance> for the given Gateway and Model so callers know expected behavior.src/proxy/types.rs (1)
7-122: Add///docs to the exported proxy types.This module now exposes a fairly large public request/response surface via
crate::proxy::types, but none of the new structs is documented. A short doc comment per exported type would make the legacy boundary much easier to consume and maintain.As per coding guidelines,
**/*.rs: Use /// for doc comments on public items in Rust.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/proxy/types.rs` around lines 7 - 122, Add concise /// doc comments for each exported public type to document the proxy request/response surface: add one-line summaries (and optional brief notes about important fields) using /// above ChatMessage, ChatCompletionResponseFormat, ChatCompletionRequest, ChatCompletionChoice, ChatCompletionUsage, ChatCompletionResponse, ChatCompletionChunkDelta, ChatCompletionChunkToolCall, ChatCompletionChunkToolCallFunction, ChatCompletionChunkChoice, and ChatCompletionChunk; keep comments short (one sentence) describing the role of the struct (e.g., "Represents a chat message with role and content"), place them immediately above the pub struct declarations, and preserve existing derives (Serialize/Deserialize/Debug/Clone/Default).
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@src/proxy/handlers/chat_completions/types.rs`:
- Around line 33-39: The current mapping in ChatCompletionError::GatewayError
(in types.rs) leaks raw upstream/transport text via err.to_string() into the
public JSON; change the produced "message" so that for server errors
(err.status_code().is_server_error() / 5xx) you return a generic message like
"Upstream provider error" and for non-server/client-safe cases only include a
sanitized/whitelisted excerpt (e.g., truncated or stripped of URLs/headers)
while preserving the existing "type" and "code" fields; update the Json
construction where ChatCompletionError::GatewayError is matched to conditionally
choose the generic message for 5xx and a sanitized err string for safe cases
(use err.status_code() to branch and a small sanitizer helper to truncate/remove
sensitive patterns).
In `@src/proxy/provider.rs`:
- Around line 23-28: The provider lookup and stored config parsing are
server-side/config errors, not client validation failures; replace the
GatewayError::Validation construction in the provider registry lookup (the let
def = gateway.registry().get(provider_name)... block) with the gateway/internal
error variant (e.g., GatewayError::Internal or the project’s internal/gateway
error variant) and do the same for the similar occurrence around the api_base
parsing (the code at the other occurrence referenced in the comment). Ensure the
new error includes the same contextual message (provider name or malformed
api_base) so logs retain useful info while returning a 5xx/internal error
classification.
---
Nitpick comments:
In `@src/gateway/traits/provider.rs`:
- Around line 324-332: Add additional unit tests for build_url on DummyProvider
to cover edge cases: assert build_url("https://example.com", "ignored") returns
"https://example.com/chat/completions" (no overlapping version), assert
build_url("https://example.com/v1/", "ignored") handles trailing slash and
returns "https://example.com/v1/chat/completions", and assert
build_url("not-a-valid-url", "ignored") exercises the fallback path behavior;
add these alongside the existing build_url_avoids_duplicate_version_prefixes
test and ensure the tests reference DummyProvider::build_url to locate the
implementation.
In `@src/proxy/hooks/observability/mod.rs`:
- Around line 61-90: Add the #[fastrace::trace] attribute to each public
observability hook to enable distributed tracing: place #[fastrace::trace]
directly above the async fn declarations for record_start_time, record_usage,
record_first_token_latency, and record_streaming_usage; ensure the fastrace
macro is in scope (import if necessary) so the attribute compiles cleanly.
In `@src/proxy/provider.rs`:
- Around line 14-21: Add missing documentation comments and fastrace decorators
to the public helper functions: annotate create_legacy_provider and
create_provider_instance with a brief triple-slash doc comment describing
purpose, parameters (e.g., model, gateway) and return value, and add the
#[fastrace::trace] attribute above each function to enable distributed tracing;
ensure the docs mention that create_legacy_provider returns a Box<dyn Provider>
created from model.provider_config and that create_provider_instance returns
Result<ProviderInstance> for the given Gateway and Model so callers know
expected behavior.
In `@src/proxy/types.rs`:
- Around line 7-122: Add concise /// doc comments for each exported public type
to document the proxy request/response surface: add one-line summaries (and
optional brief notes about important fields) using /// above ChatMessage,
ChatCompletionResponseFormat, ChatCompletionRequest, ChatCompletionChoice,
ChatCompletionUsage, ChatCompletionResponse, ChatCompletionChunkDelta,
ChatCompletionChunkToolCall, ChatCompletionChunkToolCallFunction,
ChatCompletionChunkChoice, and ChatCompletionChunk; keep comments short (one
sentence) describing the role of the struct (e.g., "Represents a chat message
with role and content"), place them immediately above the pub struct
declarations, and preserve existing derives
(Serialize/Deserialize/Debug/Clone/Default).
🪄 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: CHILL
Plan: Pro
Run ID: 6ec5da92-480e-40ed-8e96-ad18704b8348
📒 Files selected for processing (15)
src/gateway/formats/anthropic_messages.rssrc/gateway/streams/bridged.rssrc/gateway/traits/provider.rssrc/gateway/types/common.rssrc/main.rssrc/proxy/handlers/chat_completions/mod.rssrc/proxy/handlers/chat_completions/types.rssrc/proxy/handlers/embeddings/mod.rssrc/proxy/hooks/mod.rssrc/proxy/hooks/observability/mod.rssrc/proxy/hooks/rate_limit/mod.rssrc/proxy/mod.rssrc/proxy/provider.rssrc/proxy/types.rstests/package.json
💤 Files with no reviewable changes (1)
- src/proxy/hooks/mod.rs
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (2)
src/gateway/error.rs (1)
168-179: Usepretty_assertions::assert_eqin these new unit tests.These assertions are in
src/**, so the repo convention is to import the pretty-assertions macro for clearer diffs on failures.As per coding guidelines, "Use pretty_assertions::assert_eq and assert_matches::assert_matches for better test output".
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/gateway/error.rs` around lines 168 - 179, Replace the standard assert_eq usage in the new unit tests with the repo-preferred pretty_assertions::assert_eq to get better diffs: either add a use pretty_assertions::assert_eq; at the top of the test module or change the two assert_eq calls in the tests (including the one in internal_error_is_not_retryable_and_maps_to_500 and the preceding test that asserts GatewayError::Internal(...).to_string()) to call pretty_assertions::assert_eq!(...) so the tests follow the project's testing convention.src/gateway/traits/provider.rs (1)
226-229: Switch these tests topretty_assertions::assert_eq.The new URL-construction coverage is useful; it should use the repo’s preferred assertion macro so failures show an actual diff.
As per coding guidelines, "Use pretty_assertions::assert_eq and assert_matches::assert_matches for better test output".
Also applies to: 355-392
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/gateway/traits/provider.rs` around lines 226 - 229, Replace the standard assert_eq usages in the tests covering the URL-construction sections with the repo-preferred pretty_assertions::assert_eq: add use pretty_assertions::assert_eq; inside the test module (near the existing use statements like std::borrow::Cow, http::HeaderMap, serde_json::json) and change the assertion calls in the blocks around the provided ranges (the URL-construction tests and the tests at 355-392) from assert_eq!(...) to assert_eq!(...) from pretty_assertions so test failures produce diffs; ensure the import is only in cfg(test) scope if the module is shared with non-test code.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@src/proxy/handlers/chat_completions/types.rs`:
- Around line 35-43: The match arm in the error mapping in types.rs is
incorrectly classifying GatewayError::Internal as a provider_error; update the
match in the error handling (the tuple assignment for (message, error_type,
code) where GatewayError::Provider, GatewayError::Http, GatewayError::Stream are
grouped) to remove GatewayError::Internal from that provider_error branch and
add a separate branch for GatewayError::Internal that returns a distinct
message/type/code (e.g., indicating an internal/gateway configuration error) so
create_provider_instance() errors are not reported as provider errors; locate
the mapping near the error match and adjust the arms accordingly.
In `@src/proxy/provider.rs`:
- Around line 52-55: The parsed api_base value currently only checks URL syntax
via Url::parse in the base_url_override construction; add a scheme validation
after parsing to reject non-HTTP schemes by returning a GatewayError::Internal
(e.g., check url.scheme() and ensure it is "http" or "https"), so when
Some(api_base) is provided the code parses the URL and then validates scheme
(rejecting file:, mailto:, etc.) before wrapping it in Some(...).
---
Nitpick comments:
In `@src/gateway/error.rs`:
- Around line 168-179: Replace the standard assert_eq usage in the new unit
tests with the repo-preferred pretty_assertions::assert_eq to get better diffs:
either add a use pretty_assertions::assert_eq; at the top of the test module or
change the two assert_eq calls in the tests (including the one in
internal_error_is_not_retryable_and_maps_to_500 and the preceding test that
asserts GatewayError::Internal(...).to_string()) to call
pretty_assertions::assert_eq!(...) so the tests follow the project's testing
convention.
In `@src/gateway/traits/provider.rs`:
- Around line 226-229: Replace the standard assert_eq usages in the tests
covering the URL-construction sections with the repo-preferred
pretty_assertions::assert_eq: add use pretty_assertions::assert_eq; inside the
test module (near the existing use statements like std::borrow::Cow,
http::HeaderMap, serde_json::json) and change the assertion calls in the blocks
around the provided ranges (the URL-construction tests and the tests at 355-392)
from assert_eq!(...) to assert_eq!(...) from pretty_assertions so test failures
produce diffs; ensure the import is only in cfg(test) scope if the module is
shared with non-test code.
🪄 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: CHILL
Plan: Pro
Run ID: 0cdb7906-a566-4119-a657-9600076b6b55
📒 Files selected for processing (5)
src/gateway/error.rssrc/gateway/traits/provider.rssrc/proxy/handlers/chat_completions/types.rssrc/proxy/provider.rssrc/proxy/types.rs
🚧 Files skipped from review as they are similar to previous changes (1)
- src/proxy/types.rs
There was a problem hiding this comment.
🧹 Nitpick comments (1)
src/proxy/provider.rs (1)
16-18: Consider unifying config normalization between legacy and gateway paths.
create_provider_instanceuses centralized auth/base-url parsing + validation, whilecreate_legacy_providerdelegates directly tocreate_provider. A shared normalization path would reduce long-term behavior drift between embeddings/chat integrations.Also applies to: 44-72
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/proxy/provider.rs` around lines 16 - 18, The legacy path currently calls create_provider(&model.provider_config) directly causing divergence from create_provider_instance's centralized auth/base-url parsing and validation; update create_legacy_provider to reuse the same normalization/validation logic used by create_provider_instance (or call a shared helper that both create_provider_instance and create_legacy_provider use) so provider config is parsed, validated and normalized consistently; specifically, refactor or call the common normalization routine used by create_provider_instance and then pass the normalized config into create_provider (or into the unified provider factory) to avoid behavior drift between embeddings/chat integrations.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Nitpick comments:
In `@src/proxy/provider.rs`:
- Around line 16-18: The legacy path currently calls
create_provider(&model.provider_config) directly causing divergence from
create_provider_instance's centralized auth/base-url parsing and validation;
update create_legacy_provider to reuse the same normalization/validation logic
used by create_provider_instance (or call a shared helper that both
create_provider_instance and create_legacy_provider use) so provider config is
parsed, validated and normalized consistently; specifically, refactor or call
the common normalization routine used by create_provider_instance and then pass
the normalized config into create_provider (or into the unified provider
factory) to avoid behavior drift between embeddings/chat integrations.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 9fe94502-a13b-4eee-863d-bb3a9b0f29a4
📒 Files selected for processing (2)
src/proxy/handlers/chat_completions/types.rssrc/proxy/provider.rs
🚧 Files skipped from review as they are similar to previous changes (1)
- src/proxy/handlers/chat_completions/types.rs
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary by CodeRabbit
Bug Fixes
Improvements
Tests