fix(protocol): check_function_call_usage raises raw KeyError on malformed tool_choice dict

[chuenchen309]

Problem

ChatCompletionRequest.tool_choice is typed as
Optional[Union[Literal["none", "auto"], Dict]] — the Dict case has no
nested field-level schema, so pydantic accepts any dict shape at request
parse time. check_function_call_usage() then indexes into that dict
directly:

if isinstance(self.tool_choice, dict):
    if self.tool_choice["type"] != "function":
        raise BadRequestError("Only 'function' tool choice is supported")

    if len(self.tool_choice["function"]) > 1:
        raise BadRequestError("Only one tool is supported when tool_choice is specified")

    for tool in self.tools:
        if tool.function.name == self.tool_choice["function"]["name"]:
            ...

A request body like tool_choice: {"type": "function", "function": {}}
(missing "name") or tool_choice: {"function": {"name": "foo"}} (missing
"type") is schema-valid per the type hint, but crashes with a raw
KeyError instead of the intended BadRequestError.

Why it matters

Every public MLCEngine/AsyncMLCEngine chat-completion method documents:

Raises: e : BadRequestError — BadRequestError is raised when the request
is invalid.

check_function_call_usage() is called directly (no try/except) from
engine_base.process_chat_completion_request, which every one of those
documented methods calls. A caller written to the documented contract —
catching BadRequestError to produce a graceful 400-style error response —
lets this KeyError propagate uncaught instead.

Repro (bare, before fix)

req = ChatCompletionRequest(
    messages=[{"role": "user", "content": "hi"}],
    model="test",
    tools=[{"type": "function", "function": {"name": "foo", "description": "d", "parameters": {}}}],
    tool_choice={"type": "function", "function": {}},  # missing "name"
)
req.check_function_call_usage(conv_template)
# KeyError: 'name'   (expected: BadRequestError)

Fix

Use .get("type") for the type check, and validate the function
sub-object (must be a dict containing "name") before indexing into it —
raising the existing BadRequestError with a descriptive message in both
new-guard cases. The valid-request code path is unchanged.

Testing

Added tests/python/protocol/test_openai_api_protocol.py
(pytest.mark.unittest — no GPU/TVM build required) covering: missing
"name", missing "type", and the existing valid-tool_choice path.

This local environment doesn't have TVM built (a full native build wasn't
feasible for verifying a pure-Python protocol-parsing fix), so I
independently verified red→green with an importlib-based harness that
loads mlc_llm/protocol/*.py directly (pydantic/shortuuid/fastapi
only, stubbing out the parent mlc_llm package to avoid its TVM import) —
functionally equivalent to the pytest cases above:

• unfixed code: both malformed-tool_choice cases raised a raw KeyError
• fixed code: both raise the documented BadRequestError; the valid-choice
  path still correctly sets conv_template.use_function_calling = True

ruff check and ruff format --check: clean.

Duplicate-work check

Open PR #3190 (long-standing, approved, unmerged — created 2025-03-26)
touches this same function while adding structural-tag-based tool call
support, but its diff only removes two blank lines around
check_function_call_usage; it doesn't add the missing-key guards this PR
adds. No functional overlap — just a possible trivial rebase if #3190
lands first.

---

Disclosure: I used AI assistance (Claude) to find this bug and draft the
fix and tests. I independently traced the call chain from the documented
BadRequestError contract down to the unguarded dict access, reproduced
both crash cases and the valid-path behavior against the actual module
code (via the importlib harness described above, since a full TVM build
wasn't available), and take responsibility for the change's correctness.

[gemini-code-assist]

Code Review

This pull request improves the validation of tool_choice in the OpenAI API protocol implementation by safely retrieving fields and ensuring tool_choice.function is a dictionary containing a 'name' field, preventing potential KeyError exceptions. It also adds unit tests to verify this behavior. The reviewer pointed out that checking len(tool_choice_function) > 1 is redundant and potentially buggy, as it counts the number of keys in the dictionary rather than the number of tools, and suggested removing this check.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

The check len(tool_choice_function) > 1 counts the number of keys in the tool_choice_function dictionary (e.g., {"name": "foo"} has a length of 1). It does not check the number of tools specified, as tool_choice_function is already a single dictionary representing one function choice. If a client sends any extra or custom keys in the function object, this check will raise a BadRequestError with a misleading message ("Only one tool is supported when tool_choice is specified"). Since tool_choice as a dictionary can only specify a single function anyway, this check is redundant and should be removed.

Suggested change

	tool_choice_function = self.tool_choice.get("function")
	if not isinstance(tool_choice_function, dict) or "name" not in tool_choice_function:
	raise BadRequestError("tool_choice.function must be an object with a 'name' field")
	if len(tool_choice_function) > 1:
	raise BadRequestError("Only one tool is supported when tool_choice is specified")
	tool_choice_function = self.tool_choice.get("function")
	if not isinstance(tool_choice_function, dict) or "name" not in tool_choice_function:
	raise BadRequestError("tool_choice.function must be an object with a 'name' field")