From bdecd5497c4b29127326cfe660588d9538b293a1 Mon Sep 17 00:00:00 2001
From: TensorNull <tensor.null@gmail.com>
Date: Wed, 3 Jun 2026 21:03:39 +0800
Subject: [PATCH] fix: address cometapi review comments

---
 litellm/images/main.py                                |  2 +-
 litellm/llms/cometapi/common_utils.py                 |  3 +++
 .../llms/cometapi/test_cometapi_endpoints.py          | 11 +++++++++--
 3 files changed, 13 insertions(+), 3 deletions(-)

diff --git a/litellm/images/main.py b/litellm/images/main.py
index 1daa55a182a..9920f3ce1bf 100644
--- a/litellm/images/main.py
+++ b/litellm/images/main.py
@@ -442,8 +442,8 @@ def image_generation(  # noqa: PLR0915
             if custom_llm_provider == litellm.LlmProviders.COMETAPI:
                 _api_base = api_base
                 api_key = api_key or dynamic_api_key or litellm.cometapi_key
+                litellm_params_dict.pop("api_key", None)
             litellm_params_dict["api_base"] = _api_base
-            litellm_params_dict["api_key"] = api_key
 
             return llm_http_handler.image_generation_handler(
                 api_key=api_key,
diff --git a/litellm/llms/cometapi/common_utils.py b/litellm/llms/cometapi/common_utils.py
index 28bcc1307e6..eb2b8a35567 100644
--- a/litellm/llms/cometapi/common_utils.py
+++ b/litellm/llms/cometapi/common_utils.py
@@ -25,6 +25,9 @@ def get_cometapi_api_base(api_base: Optional[str] = None) -> str:
 def get_cometapi_complete_url(api_base: Optional[str], endpoint: str) -> str:
     base_url = get_cometapi_api_base(api_base).rstrip("/")
     normalized_endpoint = endpoint.strip("/")
+    parsed_endpoint = urlsplit(normalized_endpoint)
+    if not normalized_endpoint or parsed_endpoint.query or parsed_endpoint.fragment:
+        raise ValueError("CometAPI endpoint must be a non-empty path")
     parsed_base_url = urlsplit(base_url)
     path_segments = [segment for segment in parsed_base_url.path.split("/") if segment]
     endpoint_segments = normalized_endpoint.split("/")
diff --git a/tests/test_litellm/llms/cometapi/test_cometapi_endpoints.py b/tests/test_litellm/llms/cometapi/test_cometapi_endpoints.py
index ab7e681d476..2d01cb41755 100644
--- a/tests/test_litellm/llms/cometapi/test_cometapi_endpoints.py
+++ b/tests/test_litellm/llms/cometapi/test_cometapi_endpoints.py
@@ -121,6 +121,14 @@ def test_cometapi_complete_url_rejects_non_v1_paths(api_base, monkeypatch):
         get_cometapi_complete_url(api_base, "embeddings")
 
 
+@pytest.mark.parametrize("endpoint", ["", "/", "embeddings?foo=bar", "embeddings#frag"])
+def test_cometapi_complete_url_rejects_invalid_endpoints(endpoint, monkeypatch):
+    _clear_cometapi_env(monkeypatch)
+
+    with pytest.raises(ValueError, match="CometAPI endpoint must be a non-empty path"):
+        get_cometapi_complete_url("https://api.cometapi.com/v1", endpoint)
+
+
 def test_cometapi_embedding_uses_api_key_alias(monkeypatch):
     _clear_cometapi_env(monkeypatch)
     monkeypatch.setenv("COMETAPI_API_KEY", "comet-env-key")
@@ -253,6 +261,7 @@ def test_cometapi_image_generation_maps_new_openai_params(monkeypatch):
     assert call_kwargs["api_key"] == "comet-explicit-key"
     assert call_kwargs["custom_llm_provider"] == "cometapi"
     assert call_kwargs["litellm_params"]["api_base"] is None
+    assert "api_key" not in call_kwargs["litellm_params"]
     assert (
         call_kwargs["image_generation_optional_request_params"]["output_compression"]
         == 70
@@ -510,8 +519,6 @@ def test_provider_endpoint_matrix_only_updates_cometapi():
         "moderations",
     ):
         assert providers["cometapi"]["endpoints"][endpoint] is True
-        assert providers["a2a"]["endpoints"][endpoint] is False
-        assert providers["bedrock"]["endpoints"][endpoint] is False
 
 
 class _ModerationResponse: