Fix MCP fn_wrapper to pass None for optional params with UNSET defaults by yangbaechu · Pull Request #21051 · mlflow/mlflow

yangbaechu · 2026-02-21T15:31:40Z

Related Issues/PRs

What changes are proposed in this pull request?

This PR fixes mlflow/mcp/server.py:fn_wrapper to handle Click 8.3.0+ click.core.Sentinel.UNSET defaults without breaking callback invocation.

When a parameter is missing from kwargs:

If param.default is Sentinel.UNSET, this PR sets kwargs[param.name] = None.
Otherwise, it sets kwargs[param.name] = param.default (existing behavior).

This avoids passing UNSET into callbacks while ensuring omitted optional params are still present for command.callback(**kwargs).

How is this PR tested?

Manual tests

Verified on Click 8.3.1.

Python-level

from mlflow.mcp.server import fn_wrapper
from mlflow.experiments import commands as experiments_cli

wrapped = fn_wrapper(experiments_cli.commands["get"])
wrapped(experiment_id="4", output="json")

Before this PR: raises TypeError: get_experiment() missing 1 required positional argument: 'experiment_name'
With this PR: succeeds and returns the expected JSON response.

End-to-end MCP tool calls

get_experiment succeeds when exactly one of experiment_id / experiment_name is provided, and fails with the expected Click validation error when both/neither are provided.
search_traces, list_scorers, and create_run behave as expected with optional params omitted / flags set / defaults applied.

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.

Fixes MCP tool invocation on Click 8.3.0+ when optional parameters are omitted by treating Sentinel.UNSET defaults as None in fn_wrapper.

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)

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

github-actions · 2026-02-21T15:31:52Z

🛠 DevTools 🛠

Install mlflow from this PR

# mlflow
pip install git+https://github.com/mlflow/mlflow.git@refs/pull/21051/merge
# mlflow-skinny
pip install git+https://github.com/mlflow/mlflow.git@refs/pull/21051/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/21051/merge

daniellok-db

lg but let's add a test case

daniellok-db · 2026-02-23T07:19:30Z

mlflow/mcp/server.py

-                    kwargs[param.name] = param.default
+                if param.name not in kwargs:
+                    if param.default is click_unset:
+                        kwargs[param.name] = None


what's the difference here? previously would kwargs simply not contain param.name? what error does that cause?

could you include a test case that demonstrates where it was failing before?

Yes — kwargs simply didn't contain the param, so command.callback(**kwargs) raises e.g. TypeError: get_experiment() missing 1 required positional argument: 'experiment_name'.

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

yangbaechu · 2026-02-23T08:42:49Z

Added a unit test that simulates Sentinel.UNSET via monkeypatch to verify the fix regardless of Click version.

daniellok-db · 2026-02-23T07:19:56Z

mlflow/mcp/server.py

-                    kwargs[param.name] = param.default
+                if param.name not in kwargs:
+                    if param.default is click_unset:
+                        kwargs[param.name] = None


could you include a test case that demonstrates where it was failing before?

github-actions · 2026-02-23T10:04:45Z

Documentation preview for 366ada2 is available at:

https://pr-21051--mlflow-docs-preview.netlify.app/docs/latest/

More info

Ignore this comment if this PR does not change the documentation.
The preview is updated when a new commit is pushed to this PR.
This comment was created by this workflow run.
The documentation was built by this workflow run.

…ts (mlflow#21051) Signed-off-by: Yang Geonhee <abab_1212@naver.com>

…ts (#21051) Signed-off-by: Yang Geonhee <abab_1212@naver.com>

Fix MCP fn_wrapper to pass None for optional params with UNSET defaults

0be5cf3

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

github-actions bot added community Community/external contribution size/XS area/tracing MLflow Tracing and its integrations rn/bug-fix Mention under Bug Fixes in Changelogs. labels Feb 21, 2026

github-actions bot added the v3.10.1 label Feb 21, 2026

daniellok-db approved these changes Feb 23, 2026

View reviewed changes

daniellok-db reviewed Feb 23, 2026

View reviewed changes

daniellok-db self-requested a review February 23, 2026 07:20

daniellok-db self-assigned this Feb 23, 2026

daniellok-db mentioned this pull request Feb 23, 2026

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

Closed

Replace integration test with unit test for UNSET default handling

366ada2

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

daniellok-db approved these changes Feb 23, 2026

View reviewed changes

daniellok-db added this pull request to the merge queue Feb 24, 2026

Merged via the queue into mlflow:master with commit 3e31dbd Feb 24, 2026
51 of 52 checks passed

github-actions bot added size/S Small PR (10-49 LoC) and removed size/XS labels Feb 24, 2026

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

Fix MCP fn_wrapper to pass None for optional params with UNSET defaul…

2b5c392

…ts (mlflow#21051) 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 to pass None for optional params with UNSET defaul…

581876f

…ts (mlflow#21051) 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 to pass None for optional params with UNSET defaul…

5ed267b

…ts (#21051) Signed-off-by: Yang Geonhee <abab_1212@naver.com>

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Fix MCP fn_wrapper to pass None for optional params with UNSET defaults#21051

Fix MCP fn_wrapper to pass None for optional params with UNSET defaults#21051
daniellok-db merged 2 commits intomlflow:masterfrom
yangbaechu:fix-mcp-click-unset-2

yangbaechu commented Feb 21, 2026

Uh oh!

github-actions bot commented Feb 21, 2026

Install mlflow from this PR

Uh oh!

daniellok-db left a comment •

edited

Loading

Uh oh!

daniellok-db Feb 23, 2026

Uh oh!

daniellok-db Feb 23, 2026

Uh oh!

yangbaechu Feb 23, 2026

Uh oh!

yangbaechu commented Feb 23, 2026

Uh oh!

daniellok-db Feb 23, 2026

Uh oh!

github-actions bot commented Feb 23, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

yangbaechu commented Feb 21, 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 21, 2026

Install mlflow from this PR

Uh oh!

daniellok-db left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

daniellok-db Feb 23, 2026

Choose a reason for hiding this comment

Uh oh!

daniellok-db Feb 23, 2026

Choose a reason for hiding this comment

Uh oh!

yangbaechu Feb 23, 2026

Choose a reason for hiding this comment

Uh oh!

yangbaechu commented Feb 23, 2026

Uh oh!

daniellok-db Feb 23, 2026

Choose a reason for hiding this comment

Uh oh!

github-actions bot commented Feb 23, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

daniellok-db left a comment •

edited

Loading