Conversation
The OpenAI Responses API uses a flat function tool_choice shape
`{"type":"function","name":"..."}`, while Chat Completions uses a
nested shape `{"type":"function","function":{"name":"..."}}`. SMG's
shared `ToolChoice` enum modeled only the nested shape, so the Responses
endpoint deserializer failed validation on spec-correct flat input
with `data did not match any variant of untagged enum ToolChoice`,
returning HTTP 400 before any router dispatch.
Add a Responses-specific input variant that accepts the flat shape and
normalizes it into the existing internal `ToolChoice::Function` so
downstream Harmony preparation, validation, and dispatch are unchanged.
Mirror the flat shape on the way out: serialize Responses request,
response builder output, and streaming `response.completed` events with
the Responses wire format. Chat Completions backcompat is preserved.
Test plan
- `cargo +nightly fmt -- --check`
- `cargo clippy --all-targets --all-features -- -D warnings`
- `cargo test`
New regression coverage in `model_gateway/tests/spec/responses.rs` and
`model_gateway/tests/api/api_endpoints_test.rs` covers the flat-input
deserialize path, the nested-input backcompat path, response-builder
serialization, and streaming completed-event shape.
Signed-off-by: Aniruddh Krovvidi <aniruddh.krovvidi@oracle.com>
Cleared pre-existing `cargo clippy --all-targets --all-features -- -D warnings`
failures encountered while validating an unrelated fix. Mechanical
rewrites only; no behavior change.
- collapsible_match / collapsible_if: convert nested `if`/`if let` to
match arms with guards (chat_template, bucket, accumulator,
parser_endpoints_test, common/mod).
- manual_div_ceil: replace `if denom > 0 { a / denom } else { fallback }`
with `a.checked_div(denom).unwrap_or(fallback)` (mcp metrics, tokenizer
fingerprint).
- explicit_into_iter_loop: drop redundant `.into_iter()` (metrics_aggregator).
- unnecessary_sort_by: replace `sort_by(|a, b| b.cmp(&a))` with
`sort_by_key(... Reverse(...))` (data_connector memory).
Also fix `model_gateway/tests/otel_tracing_test.rs`: install the OTEL
tracing layer as the default subscriber inside the test so traces are
captured locally and in CI; previously the layer was created but never
set, causing the test to depend on harness ordering.
Test plan
- `cargo +nightly fmt -- --check`
- `cargo clippy --all-targets --all-features -- -D warnings`
- `cargo test`
Signed-off-by: Aniruddh Krovvidi <aniruddh.krovvidi@oracle.com>
📝 WalkthroughWalkthroughThis PR modernizes ChangesResponses tool_choice JSON Value Refactoring
Arithmetic Safety & Code Cleanups
Sequence Diagram(s)sequenceDiagram
participant Client
participant ResponsesParser
participant ToolChoiceDeser
participant ResponseBuilder
participant ToolChoiceEmitter
Client->>ResponsesParser: POST /v1/responses with<br/>flat function tool_choice
ResponsesParser->>ToolChoiceDeser: Deserialize request.tool_choice<br/>{ "type": "function", "name": ... }
ToolChoiceDeser->>ToolChoiceDeser: Map to ToolChoice::Function
ResponsesParser->>ResponsesParser: Validate function exists in tools
ResponsesParser->>ResponseBuilder: Build response via<br/>copy_from_request()
ResponseBuilder->>ToolChoiceEmitter: Pass tool_choice Value
ToolChoiceEmitter->>ToolChoiceEmitter: Call responses_tool_choice_value()<br/>to serialize back to flat JSON
ToolChoiceEmitter->>Client: Return event with flat<br/>tool_choice JSON shape
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes Possibly related PRs
Suggested labels
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Tip 💬 Introducing Slack Agent: The best way for teams to turn conversations into code.Slack Agent is built on CodeRabbit's deep understanding of your code, so your team can collaborate across the entire SDLC without losing context.
Built for teams:
One agent for your entire SDLC. Right inside Slack. 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. Comment |
|
Hi @ankrovv, this PR has merge conflicts that must be resolved before it can be merged. Please rebase your branch: git fetch origin main
git rebase origin/main
# resolve any conflicts, then:
git push --force-with-lease |
There was a problem hiding this comment.
Actionable comments posted: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
model_gateway/src/policies/bucket.rs (1)
366-377:⚠️ Potential issue | 🟠 Major | ⚡ Quick winGuard zero-width
gapbefore boundary range math.Line 366 only avoids
worker_cnt == 0; it does not preventgap == 0(e.g., whenworker_cnt > self.l_max). Then Line 376 underflows on- 1, which can panic in debug builds.Suggested fix
- let boundary = if let Some(gap) = self.l_max.checked_div(worker_cnt) { + let boundary = if let Some(gap) = self.l_max.checked_div(worker_cnt).filter(|&g| g > 0) { self.l_max = usize::MAX; prefill_worker_urls .iter() .enumerate() .map(|(i, url)| { let min = i * gap; let max = if i == worker_cnt - 1 { self.l_max } else { - (i + 1) * gap - 1 + (i + 1).saturating_mul(gap).saturating_sub(1) }; Boundary::new(url.clone(), [min, max]) }) .collect() } else { Vec::new() };🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@model_gateway/src/policies/bucket.rs` around lines 366 - 377, The code assumes gap > 0 after calling self.l_max.checked_div(worker_cnt) but doesn't guard gap == 0, which causes (i+1)*gap - 1 to underflow; modify the boundary construction in the block using checked_div so you first check gap != 0 (e.g., match on if gap == 0 { /* handle: assign max = self.l_max for all workers or return a safe range */ } else { ... }) or replace the subtraction with a safe operation like .saturating_sub(1); update the closure that iterates prefill_worker_urls (the map with |(i, url)| { let min = i * gap; let max = if i == worker_cnt - 1 { self.l_max } else { (i + 1) * gap - 1 } }) to use one of these guards so no underflow can occur when gap == 0.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@model_gateway/tests/otel_tracing_test.rs`:
- Around line 171-173: The test installs a subscriber built from
tracing_subscriber::registry().with(otel_trace::get_otel_layer()) and sets it as
the thread-local default via tracing::subscriber::set_default, which drops all
tracing fmt output; modify the subscriber construction in otel_tracing_test.rs
(the block that calls otel_trace::get_otel_layer(),
tracing_subscriber::registry(), and tracing::subscriber::set_default) to also
attach a formatting/test-writer layer (e.g., tracing_subscriber::fmt::layer()
configured to write to the test harness) so the subscriber contains both the
OTEL layer and a fmt layer; apply the same change to the
test_grpc_trace_context_injection setup to ensure tracing::info!/warn!/error!
output is not silently discarded.
In `@model_gateway/tests/spec/responses.rs`:
- Around line 927-958: Add a round-trip serialization assertion to
test_deserialize_responses_chat_style_function_tool_choice_backcompat: after
deserializing into ResponsesRequest and validating (in the test function),
serialize the request back to JSON (e.g., via serde_json::to_value or to_string)
and assert that the output uses the flat Responses shape (matching what
test_deserialize_responses_flat_function_tool_choice verifies) — specifically
that the serialized tool_choice is the flat Function variant with a top-level
"name": "lookup_city" (not nested under "function"). This ensures the chat-style
input normalizes back to the flat output shape.
---
Outside diff comments:
In `@model_gateway/src/policies/bucket.rs`:
- Around line 366-377: The code assumes gap > 0 after calling
self.l_max.checked_div(worker_cnt) but doesn't guard gap == 0, which causes
(i+1)*gap - 1 to underflow; modify the boundary construction in the block using
checked_div so you first check gap != 0 (e.g., match on if gap == 0 { /* handle:
assign max = self.l_max for all workers or return a safe range */ } else { ...
}) or replace the subtraction with a safe operation like .saturating_sub(1);
update the closure that iterates prefill_worker_urls (the map with |(i, url)| {
let min = i * gap; let max = if i == worker_cnt - 1 { self.l_max } else { (i +
1) * gap - 1 } }) to use one of these guards so no underflow can occur when gap
== 0.
🪄 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: ed3ebb7b-84b3-496c-b4ed-537e7136daf0
📒 Files selected for processing (15)
crates/data_connector/src/memory.rscrates/mcp/src/core/metrics.rscrates/protocols/src/builders/responses/response.rscrates/protocols/src/responses.rscrates/tokenizer/src/cache/fingerprint.rscrates/tokenizer/src/chat_template.rsmodel_gateway/src/core/metrics_aggregator.rsmodel_gateway/src/policies/bucket.rsmodel_gateway/src/routers/grpc/common/responses/streaming.rsmodel_gateway/src/routers/openai/responses/accumulator.rsmodel_gateway/tests/api/api_endpoints_test.rsmodel_gateway/tests/api/parser_endpoints_test.rsmodel_gateway/tests/common/mod.rsmodel_gateway/tests/otel_tracing_test.rsmodel_gateway/tests/spec/responses.rs
| let otel_layer = otel_trace::get_otel_layer().expect("Failed to get OTEL layer"); | ||
| let subscriber = tracing_subscriber::registry().with(otel_layer); | ||
| let _subscriber_guard = tracing::subscriber::set_default(subscriber); |
There was a problem hiding this comment.
Subscriber constructed with only the OTEL layer; all tracing log output on the test thread is silently dropped
tracing_subscriber::registry().with(otel_layer) (line 172) contains no fmt/logging layer. Once set_default installs it as the thread-local default (line 173), it shadows the global subscriber created by init_logging() (which has the fmt layer) for the entire remaining test. On a #[tokio::test] current-thread runtime, every tracing::info!, tracing::warn!, and tracing::error! emitted by the router, middleware, or worker — for roughly 100 lines of async execution — is silently discarded. The test assertions are unaffected (the OTEL layer is present), but diagnosing a future test failure becomes much harder because framework instrumentation won't appear in test output.
Add a fmt test-writer layer to the subscriber so both concerns are served:
🔧 Proposed fix
let otel_layer = otel_trace::get_otel_layer().expect("Failed to get OTEL layer");
-let subscriber = tracing_subscriber::registry().with(otel_layer);
+let fmt_layer = tracing_subscriber::fmt::layer().with_test_writer();
+let subscriber = tracing_subscriber::registry().with(fmt_layer).with(otel_layer);
let _subscriber_guard = tracing::subscriber::set_default(subscriber);Note: test_grpc_trace_context_injection (lines 316–317, unchanged) has the same pattern and would benefit from the same fix.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@model_gateway/tests/otel_tracing_test.rs` around lines 171 - 173, The test
installs a subscriber built from
tracing_subscriber::registry().with(otel_trace::get_otel_layer()) and sets it as
the thread-local default via tracing::subscriber::set_default, which drops all
tracing fmt output; modify the subscriber construction in otel_tracing_test.rs
(the block that calls otel_trace::get_otel_layer(),
tracing_subscriber::registry(), and tracing::subscriber::set_default) to also
attach a formatting/test-writer layer (e.g., tracing_subscriber::fmt::layer()
configured to write to the test harness) so the subscriber contains both the
OTEL layer and a fmt layer; apply the same change to the
test_grpc_trace_context_injection setup to ensure tracing::info!/warn!/error!
output is not silently discarded.
| #[test] | ||
| fn test_deserialize_responses_chat_style_function_tool_choice_backcompat() { | ||
| let request: ResponsesRequest = serde_json::from_value(json!({ | ||
| "input": "test", | ||
| "tools": [ | ||
| { | ||
| "type": "function", | ||
| "name": "lookup_city", | ||
| "parameters": {} | ||
| } | ||
| ], | ||
| "tool_choice": { | ||
| "type": "function", | ||
| "function": { | ||
| "name": "lookup_city" | ||
| } | ||
| } | ||
| })) | ||
| .expect("Chat-style function tool_choice should remain accepted"); | ||
|
|
||
| assert!( | ||
| matches!( | ||
| request.tool_choice, | ||
| Some(ToolChoice::Function { ref function, .. }) if function.name == "lookup_city" | ||
| ), | ||
| "Chat-style tool_choice should keep using the internal function choice" | ||
| ); | ||
| assert!( | ||
| request.validate().is_ok(), | ||
| "Chat-style function tool_choice should still validate" | ||
| ); | ||
| } |
There was a problem hiding this comment.
🧹 Nitpick | 🔵 Trivial | ⚡ Quick win
test_deserialize_responses_chat_style_function_tool_choice_backcompat is missing the serialization assertion.
The PR's stated contract is that both accepted input shapes (flat Responses and nested Chat-style) mirror the flat Responses shape on output. Test 1 (test_deserialize_responses_flat_function_tool_choice, lines 853-861) verifies round-trip serialization for flat input. This backcompat test verifies deserialization and validation for Chat-style input, but it never checks that the serialized output is also in flat shape — leaving a gap in the critical "normalize → serialize" path.
✅ Proposed serialization assertion for the backcompat test
assert!(
request.validate().is_ok(),
"Chat-style function tool_choice should still validate"
);
+
+ let serialized = serde_json::to_value(&request).expect("ResponsesRequest should serialize");
+ assert_eq!(
+ serialized["tool_choice"],
+ json!({
+ "type": "function",
+ "name": "lookup_city"
+ }),
+ "Chat-style input should also serialize as flat Responses shape"
+ );
}📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| #[test] | |
| fn test_deserialize_responses_chat_style_function_tool_choice_backcompat() { | |
| let request: ResponsesRequest = serde_json::from_value(json!({ | |
| "input": "test", | |
| "tools": [ | |
| { | |
| "type": "function", | |
| "name": "lookup_city", | |
| "parameters": {} | |
| } | |
| ], | |
| "tool_choice": { | |
| "type": "function", | |
| "function": { | |
| "name": "lookup_city" | |
| } | |
| } | |
| })) | |
| .expect("Chat-style function tool_choice should remain accepted"); | |
| assert!( | |
| matches!( | |
| request.tool_choice, | |
| Some(ToolChoice::Function { ref function, .. }) if function.name == "lookup_city" | |
| ), | |
| "Chat-style tool_choice should keep using the internal function choice" | |
| ); | |
| assert!( | |
| request.validate().is_ok(), | |
| "Chat-style function tool_choice should still validate" | |
| ); | |
| } | |
| #[test] | |
| fn test_deserialize_responses_chat_style_function_tool_choice_backcompat() { | |
| let request: ResponsesRequest = serde_json::from_value(json!({ | |
| "input": "test", | |
| "tools": [ | |
| { | |
| "type": "function", | |
| "name": "lookup_city", | |
| "parameters": {} | |
| } | |
| ], | |
| "tool_choice": { | |
| "type": "function", | |
| "function": { | |
| "name": "lookup_city" | |
| } | |
| } | |
| })) | |
| .expect("Chat-style function tool_choice should remain accepted"); | |
| assert!( | |
| matches!( | |
| request.tool_choice, | |
| Some(ToolChoice::Function { ref function, .. }) if function.name == "lookup_city" | |
| ), | |
| "Chat-style tool_choice should keep using the internal function choice" | |
| ); | |
| assert!( | |
| request.validate().is_ok(), | |
| "Chat-style function tool_choice should still validate" | |
| ); | |
| let serialized = serde_json::to_value(&request).expect("ResponsesRequest should serialize"); | |
| assert_eq!( | |
| serialized["tool_choice"], | |
| json!({ | |
| "type": "function", | |
| "name": "lookup_city" | |
| }), | |
| "Chat-style input should also serialize as flat Responses shape" | |
| ); | |
| } |
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@model_gateway/tests/spec/responses.rs` around lines 927 - 958, Add a
round-trip serialization assertion to
test_deserialize_responses_chat_style_function_tool_choice_backcompat: after
deserializing into ResponsesRequest and validating (in the test function),
serialize the request back to JSON (e.g., via serde_json::to_value or to_string)
and assert that the output uses the flat Responses shape (matching what
test_deserialize_responses_flat_function_tool_choice verifies) — specifically
that the serialized tool_choice is the flat Function variant with a top-level
"name": "lookup_city" (not nested under "function"). This ensures the chat-style
input normalizes back to the flat output shape.
There was a problem hiding this comment.
Code Review
This pull request implements support for a "flat" tool choice format in the responses protocol, transitioning the tool_choice field from a String to a serde_json::Value to accommodate more complex structures. It includes custom serialization and deserialization logic to maintain backward compatibility with existing string and chat-style formats. Beyond these protocol changes, the PR performs several code quality improvements, such as replacing manual division checks with checked_div, utilizing sort_by_key for more idiomatic sorting, and refactoring match statements to use if guards. I have no feedback to provide.
|
Fixed by #1276 |
Description
Problem
POST /v1/responsesrejects requests that use the OpenAI-spec-correct flat functiontool_choiceshape with HTTP 400tool_choice: data did not match any variant of untagged enum ToolChoice.The OpenAI Responses API uses a flat shape:
{ "tool_choice": { "type": "function", "name": "lookup_city" } }while Chat Completions uses a nested shape:
{ "tool_choice": { "type": "function", "function": { "name": "lookup_city" } } }ResponsesRequest.tool_choicereuses the sharedToolChoiceenum which models only the nested Chat shape, so serde rejects the spec-correct flat input at the SMG entrypoint, before any router dispatch — affecting both gRPC- and HTTP-backend routing modes equally.Solution
Add a Responses-specific input variant that accepts the flat shape and normalize it into the existing internal
ToolChoice::Functionafter deserialization. Downstream Harmony preparation, validation, and dispatch continue to operate on the existing internal shape unchanged. The Chat Completions nested shape remains accepted via the same input variant, preserving backcompat.Changes
crates/protocols/src/responses.rs— Responses-specific tool_choice input variant accepting flat shape; normalization to internalToolChoice.crates/protocols/src/builders/responses/response.rs— flat-shape serialization for Responses output.model_gateway/src/routers/grpc/common/responses/streaming.rs— flat-shape serialization for streamingresponse.completed.model_gateway/tests/spec/responses.rs— spec coverage for flat-input deserialize, nested-input backcompat, and roundtrip.model_gateway/tests/api/api_endpoints_test.rs— endpoint guard exercising the flat-input path end-to-end.Test Plan
All three pass locally. The new spec and endpoint tests cover:
response.completedall serialize the flat shapeChecklist
cargo +nightly fmtpassescargo clippy --all-targets --all-features -- -D warningspassesSummary by CodeRabbit
New Features
Bug Fixes
Tests