Fix MCP fn_wrapper handling of Click UNSET defaults (#20953) by yangbaechu · Pull Request #20962 · mlflow/mlflow

yangbaechu · 2026-02-18T10:22:37Z

Related Issues/PRs

Resolve #20953

What changes are proposed in this pull request?

This PR fixes MCP fn_wrapper handling for Click 8.3+ by skipping click.core.UNSET when filling missing optional params, so callbacks receive their Python defaults (e.g., None) instead of Sentinel.UNSET. It also adds a regression test in tests/mcp/test_server.py for the missing-optional-arg path.

How is this PR tested?

New unit/integration tests
Manual tests

Manual reproduce run:

from mlflow.mcp.server import fn_wrapper                                                                                             
from mlflow.cli.traces import commands as traces_cli                                                                                 
                                                                                                                                     
wrapped = fn_wrapper(traces_cli.commands["search"])                                                                                  
wrapped(experiment_id="0")

Result: no AttributeError, command runs successfully.

Does this PR require documentation update?

No. You can skip the rest of this section.

Does this PR require updating the MLflow Skills repository?

No. You can skip the rest of this section.

Release Notes

Is this a user-facing change?

Yes. Give a description of this change to be included in the release notes for MLflow users.

MLflow MCP no longer passes Click 8.3+ Sentinel.UNSET for missing optional tool arguments, preventing callback crashes such as 'Sentinel' object has no attribute 'split'.

What component(s), interfaces, languages, and integrations does this PR affect?

Components

area/tracing: MLflow Tracing features, tracing APIs, and LLM tracing functionality

How should the PR be classified in the release notes? Choose one:

rn/bug-fix - A user-facing bug fix worth mentioning in the release notes

Should this PR be included in the next patch release?

Yes should be selected for bug fixes, documentation updates, and other small changes. No should be selected for new features and larger changes. If you're unsure about the release classification of this PR, leave this unchecked to let the maintainers decide.

What is a minor/patch release?

Minor release: a release that increments the second part of the version number (e.g., 1.2.0 -> 1.3.0).
Bug fixes, doc updates and new features usually go into minor releases.
Patch release: a release that increments the third part of the version number (e.g., 1.2.0 -> 1.2.1).
Bug fixes and doc updates usually go into patch releases.

Yes (this PR will be cherry-picked and included in the next patch release)
No (this PR will be included in the next minor release)

github-actions · 2026-02-18T10:22:49Z

🛠 DevTools 🛠

Install mlflow from this PR

# mlflow
pip install git+https://github.com/mlflow/mlflow.git@refs/pull/20962/merge
# mlflow-skinny
pip install git+https://github.com/mlflow/mlflow.git@refs/pull/20962/merge#subdirectory=libs/skinny

For Databricks, use the following command:

%sh curl -LsSf https://raw.githubusercontent.com/mlflow/mlflow/HEAD/dev/install-skinny.sh | sh -s pull/20962/merge

github-actions · 2026-02-18T10:22:50Z

@yangbaechu Thank you for the contribution! Could you fix the following issue(s)?

⚠ DCO check

The DCO check failed. Please sign off your commit(s) by following the instructions here. See https://github.com/mlflow/mlflow/blob/master/CONTRIBUTING.md#sign-your-work for more details.

TomeHirata · 2026-02-20T02:02:21Z

tests/mcp/test_server.py

@@ -0,0 +1,20 @@
+from unittest.mock import Mock


Can we test in test_mcp?

Done — moved to test_mcp.py

TomeHirata · 2026-02-20T02:02:56Z

mlflow/mcp/server.py


 def fn_wrapper(command: click.Command) -> Callable[..., str]:
    def wrapper(**kwargs: Any) -> str:
+        click_unset = getattr(click.core, "UNSET", object())


Is it intentional to use object() as a default?

Yes — it's a fallback sentinel for click < 8.3.0 where click.core.UNSET doesn't exist. object() ensures it never matches any param.default via is, so the new skip logic only activates on click >= 8.3.0.

Signed-off-by: Yang Geonhee <abab_1212@naver.com>

yangbaechu · 2026-02-20T05:43:09Z

I moved test_server into test_mcp and updated DCO sign-off to use my legal name and force-pushed to align the PR branch history.

daniellok-db · 2026-02-20T06:03:47Z

i'm removing the v3.10 label as the release is happening within the next 30 mins or so. it will be included in the next release which shouldn't be too far away

TomeHirata · 2026-02-20T06:13:15Z

mlflow/mcp/server.py

@@ -93,6 +95,8 @@ def wrapper(**kwargs: Any) -> str:
            # Fill in defaults for missing optional arguments
            for param in command.params:
                if param.name not in kwargs:


Can we combine the if condition?

Addressed in 56caaf6.

TomeHirata · 2026-02-20T06:13:30Z

tests/mcp/test_mcp.py

+
+@pytest.mark.asyncio
+async def test_call_tool_with_optional_params_omitted(client: Client):
+    """Regression test for click >= 8.3.0 Sentinel.UNSET handling (#20953)."""


Let's remove this docstring

Addressed in 56caaf6

Signed-off-by: Yang Geonhee <abab_1212@naver.com>

TomeHirata

LGTM

…lflow#20962) Signed-off-by: Yang Geonhee <abab_1212@naver.com>

Signed-off-by: Yang Geonhee <abab_1212@naver.com>

github-actions bot added the community Community/external contribution label Feb 18, 2026

github-actions bot added the size/S Small PR (10-49 LoC) label Feb 18, 2026

github-actions bot added v3.10.0 area/tracing MLflow Tracing and its integrations rn/bug-fix Mention under Bug Fixes in Changelogs. labels Feb 18, 2026

yangbaechu force-pushed the fix-mcp-click-unset-20953 branch from 59addcd to b35e062 Compare February 18, 2026 10:26

TomeHirata reviewed Feb 20, 2026

View reviewed changes

TomeHirata removed the v3.10.0 label Feb 20, 2026

github-actions bot assigned TomeHirata Feb 20, 2026

yangbaechu added 2 commits February 20, 2026 13:53

Fix MCP fn_wrapper handling of Click UNSET defaults (mlflow#20953)

ab534f3

Signed-off-by: Yang Geonhee <abab_1212@naver.com>

Move regression test from test_server.py to test_mcp.py

238a7ac

Signed-off-by: Yang Geonhee <abab_1212@naver.com>

yangbaechu force-pushed the fix-mcp-click-unset-20953 branch from b35e062 to 238a7ac Compare February 20, 2026 05:36

github-actions bot added the v3.10.0 label Feb 20, 2026

daniellok-db removed the v3.10.0 label Feb 20, 2026

TomeHirata reviewed Feb 20, 2026

View reviewed changes

github-actions bot assigned daniellok-db Feb 20, 2026

Address review comments: combine if condition and remove docstring

56caaf6

Signed-off-by: Yang Geonhee <abab_1212@naver.com>

github-actions bot added the v3.10.0 label Feb 20, 2026

TomeHirata approved these changes Feb 20, 2026

View reviewed changes

TomeHirata enabled auto-merge February 20, 2026 07:09

TomeHirata removed the v3.10.0 label Feb 20, 2026

TomeHirata added this pull request to the merge queue Feb 20, 2026

Merged via the queue into mlflow:master with commit fce7f83 Feb 20, 2026
52 of 54 checks passed

yangbaechu mentioned this pull request Feb 21, 2026

[Bug] MCP fn_wrapper incorrectly handles Click 8.3.0+ Sentinel.UNSET for omitted optional params #21050

Closed

11 tasks

paipeline mentioned this pull request Feb 22, 2026

Fix MCP fn_wrapper handling of Click 8.3.0+ Sentinel.UNSET parameters (#21050) #21055

Closed

daniellok-db added the v3.10.1 label Mar 5, 2026

daniellok-db pushed a commit to daniellok-db/mlflow that referenced this pull request Mar 5, 2026

Fix MCP fn_wrapper handling of Click UNSET defaults (mlflow#20953) (m…

f5ea06c

…lflow#20962) Signed-off-by: Yang Geonhee <abab_1212@naver.com>

daniellok-db pushed a commit to daniellok-db/mlflow that referenced this pull request Mar 5, 2026

Fix MCP fn_wrapper handling of Click UNSET defaults (mlflow#20953) (m…

574b959

…lflow#20962) Signed-off-by: Yang Geonhee <abab_1212@naver.com>

daniellok-db pushed a commit that referenced this pull request Mar 5, 2026

Fix MCP fn_wrapper handling of Click UNSET defaults (#20953) (#20962)

cb5c256

Signed-off-by: Yang Geonhee <abab_1212@naver.com>

Conversation

yangbaechu commented Feb 18, 2026

Related Issues/PRs

What changes are proposed in this pull request?

How is this PR tested?

Does this PR require documentation update?

Does this PR require updating the MLflow Skills repository?

Release Notes

Is this a user-facing change?

What component(s), interfaces, languages, and integrations does this PR affect?

How should the PR be classified in the release notes? Choose one:

Should this PR be included in the next patch release?

Uh oh!

github-actions bot commented Feb 18, 2026

Install mlflow from this PR

Uh oh!

github-actions bot commented Feb 18, 2026

⚠ DCO check

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

yangbaechu commented Feb 20, 2026

Uh oh!

daniellok-db commented Feb 20, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

TomeHirata left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants