kaxil commented on code in PR #74451:
URL: https://github.com/apache/airflow/pull/74451#discussion_r4218124576
##########
providers/common/ai/src/airflow/providers/common/ai/hooks/mcp.py:
##########
@@ -239,7 +239,23 @@ def get_conn(self) -> Any:
args = extra.get("args", [])
if isinstance(args, str):
args = [args]
+ if not isinstance(args, list) or not all(isinstance(arg, str) for
arg in args):
+ raise ValueError(
+ f"'args' in extra for connection {self.mcp_conn_id!r} must
be a string or a "
+ f"list of strings, got {args!r}."
+ )
timeout = extra.get("timeout", 10)
+ # ``None`` and ``0`` are kept valid: fastmcp's
``normalize_timeout_to_seconds``
+ # treats both as "no init timeout".
Review Comment:
This holds for `0` but not quite for `None`. `MCPToolset` passes
`init_timeout=None` through unchanged, and fastmcp's `Client.__init__` replaces
`None` with `fastmcp.settings.client_init_timeout` before normalizing. So on a
worker with `FASTMCP_CLIENT_INIT_TIMEOUT` set, a `null` timeout still times
out. Mapping `None` to `0` before building the toolset would make this comment
and the `mcp.rst` line true in every case.
##########
providers/common/ai/tests/unit/common/ai/hooks/test_mcp.py:
##########
@@ -207,6 +209,54 @@ def test_stdio_without_command_raises(self):
with pytest.raises(ValueError, match="requires 'command'"):
hook.get_conn()
+ def test_missing_mcp_dependency_raises_optional_feature_exception(self):
+ hook = MCPHook(mcp_conn_id="test_conn")
+ with patch.dict(sys.modules, {"fastmcp.client.transports": None}):
+ with pytest.raises(AirflowOptionalProviderFeatureException,
match="pydantic-ai-slim"):
+ hook.get_conn()
+
+ @pytest.mark.parametrize("bad_args", [{"a": 1}, ["-m", 1], 5],
ids=["dict", "non_string_item", "int"])
+ def test_stdio_invalid_args_raises(self, bad_args):
+ hook = MCPHook(mcp_conn_id="test_conn")
+ conn = Connection(
+ conn_id="test_conn",
+ conn_type="mcp",
+ extra=json.dumps({"transport": "stdio", "command": "uvx", "args":
bad_args}),
+ )
+ with patch.object(hook, "get_connection", return_value=conn):
+ with pytest.raises(ValueError, match="'args' in extra"):
+ hook.get_conn()
+
+ @pytest.mark.parametrize(
+ "bad_timeout",
+ ["30", -5, -0.5, float("nan"), float("inf"), True, [10]],
+ ids=["str", "neg", "neg_float", "nan", "inf", "bool", "list"],
+ )
+ def test_stdio_invalid_timeout_raises(self, bad_timeout):
+ hook = MCPHook(mcp_conn_id="test_conn")
+ conn = Connection(
+ conn_id="test_conn",
+ conn_type="mcp",
+ extra=json.dumps({"transport": "stdio", "command": "uvx", "args":
[], "timeout": bad_timeout}),
+ )
+ with patch.object(hook, "get_connection", return_value=conn):
+ with pytest.raises(ValueError, match="'timeout' in extra"):
+ hook.get_conn()
+
+ @patch(_MCP_TOOLSET, autospec=True)
+ @patch(_STDIO_TRANSPORT, autospec=True)
+ def test_stdio_float_zero_and_null_timeout_accepted(self,
mock_transport_cls, mock_toolset_cls):
Review Comment:
Could this be parametrized like the two tests above it? With the loop a
failure doesn't say which timeout value broke, and the mocks are shared across
iterations.
##########
providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py:
##########
@@ -412,6 +415,7 @@ def _provider_factory(pname: str) -> Any:
pname,
list(_kwargs),
)
+ self._env_auth_fallbacks.append(f"{pname} (ignored: {',
'.join(_kwargs)})")
Review Comment:
This makes `test_connection` honest, but at task runtime the swap is still
only a log warning, and the Vertex comment further down calls this path
authenticating as the wrong identity. Since this PR is about settling hook
behaviour while it can still change freely, would it be worth failing closed
here (a `ValueError` built from this same string), or making the env-var
fallback opt-in? Fine as a follow-up if you'd rather keep this PR scoped.
##########
providers/common/ai/tests/unit/common/ai/hooks/test_langchain.py:
##########
@@ -72,6 +72,11 @@ def test_embed_conn_falls_back_to_llm_conn(self):
assert hook.llm_conn_id == "my_conn"
assert hook.embed_conn_id == "my_conn"
+ def test_model_arguments_are_keyword_only(self):
+ """``llm_model`` / ``embed_model`` sit in opposite order across the
two hooks."""
Review Comment:
After this PR both hooks declare `llm_model` before `embed_model`, so this
docstring describes the old signature. Maybe say what the test pins instead
(the model arguments can't be passed positionally). Same docstring in
`test_llamaindex.py`.
##########
providers/common/ai/src/airflow/providers/common/ai/policies/retry.py:
##########
@@ -66,7 +67,7 @@
RetryPolicy,
)
except ImportError:
- raise ImportError(
+ raise AirflowOptionalProviderFeatureException(
Review Comment:
`AirflowOptionalProviderFeatureException` isn't an `ImportError` subclass,
so the `except ImportError:` guard in
`example_dags/example_llm_retry_policy.py` (line 173, "example DAG is skipped
on older versions") no longer catches this on Airflow 2.11 to 3.2. Any user Dag
that copied that guard breaks the same way. Could the example catch
`(ImportError, AirflowOptionalProviderFeatureException)`, or import
`airflow.sdk.definitions.retry_policy` first inside the `try` so older cores
still fail with `ImportError`?
##########
providers/common/ai/src/airflow/providers/common/ai/hooks/mcp.py:
##########
@@ -239,7 +239,23 @@ def get_conn(self) -> Any:
args = extra.get("args", [])
if isinstance(args, str):
args = [args]
+ if not isinstance(args, list) or not all(isinstance(arg, str) for
arg in args):
Review Comment:
Since this is now where `args` gets validated, could it also handle the two
shapes the connection form produces? `provider.yaml` declares `args` as a
nullable string and the form writes blank fields as `null`, so a stdio
connection with Arguments left empty ends up with `"args": null` and is
rejected here (it already failed before, just later in
`StdioServerParameters`). And the field description asks for a JSON array, but
typing `["mcp-run-python"]` gets wrapped as one literal argument by the
`isinstance(args, str)` branch above and passes this check. Treating `None` as
`[]` and `json.loads`-ing a string that starts with `[` would cover both. I
traced the form path in code, didn't run the UI.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
