From eef117bb790ec3ac347585b6ec90a6bdfcbd5626 Mon Sep 17 00:00:00 2001 From: fzowl Date: Fri, 27 Dec 2024 22:04:02 +0100 Subject: [PATCH 01/21] Refresh VoyageAI models and prices and context --- model_prices_and_context_window.json | 94 ++++++++++++++++++---------- 1 file changed, 62 insertions(+), 32 deletions(-) diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 5f98c0e68b0..a12d114f952 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -7566,30 +7566,6 @@ "litellm_provider": "cloudflare", "mode": "chat" }, - "voyage/voyage-01": { - "max_tokens": 4096, - "max_input_tokens": 4096, - "input_cost_per_token": 0.0000001, - "output_cost_per_token": 0.000000, - "litellm_provider": "voyage", - "mode": "embedding" - }, - "voyage/voyage-lite-01": { - "max_tokens": 4096, - "max_input_tokens": 4096, - "input_cost_per_token": 0.0000001, - "output_cost_per_token": 0.000000, - "litellm_provider": "voyage", - "mode": "embedding" - }, - "voyage/voyage-large-2": { - "max_tokens": 16000, - "max_input_tokens": 16000, - "input_cost_per_token": 0.00000012, - "output_cost_per_token": 0.000000, - "litellm_provider": "voyage", - "mode": "embedding" - }, "voyage/voyage-law-2": { "max_tokens": 16000, "max_input_tokens": 16000, @@ -7614,14 +7590,6 @@ "litellm_provider": "voyage", "mode": "embedding" }, - "voyage/voyage-lite-02-instruct": { - "max_tokens": 4000, - "max_input_tokens": 4000, - "input_cost_per_token": 0.0000001, - "output_cost_per_token": 0.000000, - "litellm_provider": "voyage", - "mode": "embedding" - }, "voyage/voyage-finance-2": { "max_tokens": 4000, "max_input_tokens": 4000, @@ -7630,6 +7598,68 @@ "litellm_provider": "voyage", "mode": "embedding" }, + "voyage/voyage-3-large": { + "max_tokens": 32000, + "max_input_tokens": 32000, + "input_cost_per_token": 0.00000018, + "output_cost_per_token": 0.000000, + "litellm_provider": "voyage", + "mode": "embedding" + }, + "voyage/voyage-3": { + "max_tokens": 32000, + "max_input_tokens": 32000, + "input_cost_per_token": 0.00000006, + "output_cost_per_token": 0.000000, + "litellm_provider": "voyage", + "mode": "embedding" + }, + "voyage/voyage-3-lite": { + "max_tokens": 32000, + "max_input_tokens": 32000, + "input_cost_per_token": 0.00000002, + "output_cost_per_token": 0.000000, + "litellm_provider": "voyage", + "mode": "embedding" + }, + "voyage/voyage-code-3": { + "max_tokens": 32000, + "max_input_tokens": 32000, + "input_cost_per_token": 0.00000018, + "output_cost_per_token": 0.000000, + "litellm_provider": "voyage", + "mode": "embedding" + }, + "voyage/voyage-multimodal-3": { + "max_tokens": 32000, + "max_input_tokens": 32000, + "input_cost_per_token": 0.00000012, + "output_cost_per_token": 0.000000, + "litellm_provider": "voyage", + "mode": "embedding" + }, + "voyage/rerank-2": { + "max_tokens": 16000, + "max_input_tokens": 16000, + "max_output_tokens": 16000, + "max_query_tokens": 16000, + "input_cost_per_token": 0.00000005, + "input_cost_per_query": 0.00000005, + "output_cost_per_token": 0.0, + "litellm_provider": "voyage", + "mode": "rerank" + }, + "voyage/rerank-2-lite": { + "max_tokens": 8000, + "max_input_tokens": 8000, + "max_output_tokens": 8000, + "max_query_tokens": 8000, + "input_cost_per_token": 0.00000002, + "input_cost_per_query": 0.00000002, + "output_cost_per_token": 0.0, + "litellm_provider": "voyage", + "mode": "rerank" + }, "databricks/databricks-meta-llama-3-1-405b-instruct": { "max_tokens": 128000, "max_input_tokens": 128000, From 6876a14838605865276ac17ea9991e1874f8d64a Mon Sep 17 00:00:00 2001 From: fzowl Date: Mon, 30 Dec 2024 11:50:37 +0100 Subject: [PATCH 02/21] Refresh VoyageAI models and prices and context --- model_prices_and_context_window.json | 32 ++++++++++++++++++++++++++++ 1 file changed, 32 insertions(+) diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 2ff5e40247d..0701b20af49 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -7613,6 +7613,38 @@ "litellm_provider": "cloudflare", "mode": "chat" }, + "voyage/voyage-01": { + "max_tokens": 4096, + "max_input_tokens": 4096, + "input_cost_per_token": 0.0000001, + "output_cost_per_token": 0.000000, + "litellm_provider": "voyage", + "mode": "embedding" + }, + "voyage/voyage-lite-01": { + "max_tokens": 4096, + "max_input_tokens": 4096, + "input_cost_per_token": 0.0000001, + "output_cost_per_token": 0.000000, + "litellm_provider": "voyage", + "mode": "embedding" + }, + "voyage/voyage-large-2": { + "max_tokens": 16000, + "max_input_tokens": 16000, + "input_cost_per_token": 0.00000012, + "output_cost_per_token": 0.000000, + "litellm_provider": "voyage", + "mode": "embedding" + }, + "voyage/voyage-lite-02-instruct": { + "max_tokens": 4000, + "max_input_tokens": 4000, + "input_cost_per_token": 0.0000001, + "output_cost_per_token": 0.000000, + "litellm_provider": "voyage", + "mode": "embedding" + }, "voyage/voyage-law-2": { "max_tokens": 16000, "max_input_tokens": 16000, From bb0f74c48f9a85874f2f104a6de0f043d4658ea6 Mon Sep 17 00:00:00 2001 From: fzowl Date: Mon, 30 Dec 2024 11:54:37 +0100 Subject: [PATCH 03/21] Refresh VoyageAI models and prices and context --- model_prices_and_context_window.json | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 0701b20af49..1a9ed533e38 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -7637,6 +7637,14 @@ "litellm_provider": "voyage", "mode": "embedding" }, + "voyage/voyage-finance-2": { + "max_tokens": 32000, + "max_input_tokens": 32000, + "input_cost_per_token": 0.00000012, + "output_cost_per_token": 0.000000, + "litellm_provider": "voyage", + "mode": "embedding" + }, "voyage/voyage-lite-02-instruct": { "max_tokens": 4000, "max_input_tokens": 4000, @@ -7669,14 +7677,6 @@ "litellm_provider": "voyage", "mode": "embedding" }, - "voyage/voyage-finance-2": { - "max_tokens": 4000, - "max_input_tokens": 4000, - "input_cost_per_token": 0.00000012, - "output_cost_per_token": 0.000000, - "litellm_provider": "voyage", - "mode": "embedding" - }, "voyage/voyage-3-large": { "max_tokens": 32000, "max_input_tokens": 32000, From 9690cb66c6e51f81db586152d822965e1db94a60 Mon Sep 17 00:00:00 2001 From: fzowl Date: Mon, 3 Feb 2025 15:21:20 +0100 Subject: [PATCH 04/21] Updating the available VoyageAI models in the docs --- docs/my-website/docs/providers/voyage.md | 27 +++++++++++++++--------- 1 file changed, 17 insertions(+), 10 deletions(-) diff --git a/docs/my-website/docs/providers/voyage.md b/docs/my-website/docs/providers/voyage.md index a56a1408ea9..6ab6b1846f5 100644 --- a/docs/my-website/docs/providers/voyage.md +++ b/docs/my-website/docs/providers/voyage.md @@ -14,7 +14,7 @@ import os os.environ['VOYAGE_API_KEY'] = "" response = embedding( - model="voyage/voyage-01", + model="voyage/voyage-3-large", input=["good morning from litellm"], ) print(response) @@ -23,13 +23,20 @@ print(response) ## Supported Models All models listed here https://docs.voyageai.com/embeddings/#models-and-specifics are supported -| Model Name | Function Call | -|--------------------------|------------------------------------------------------------------------------------------------------------------------------------------------------------------| -| voyage-2 | `embedding(model="voyage/voyage-2", input)` | -| voyage-large-2 | `embedding(model="voyage/voyage-large-2", input)` | -| voyage-law-2 | `embedding(model="voyage/voyage-law-2", input)` | -| voyage-code-2 | `embedding(model="voyage/voyage-code-2", input)` | +| Model Name | Function Call | +|-------------------------|------------------------------------------------------------| +| voyage-3-large | `embedding(model="voyage/voyage-3-large", input)` | +| voyage-3 | `embedding(model="voyage/voyage-3", input)` | +| voyage-3-lite | `embedding(model="voyage/voyage-3-lite", input)` | +| voyage-code-3 | `embedding(model="voyage/voyage-code-3", input)` | +| voyage-finance-2 | `embedding(model="voyage/voyage-finance-2", input)` | +| voyage-law-2 | `embedding(model="voyage/voyage-law-2", input)` | +| voyage-code-2 | `embedding(model="voyage/voyage-code-2", input)` | +| voyage-multilingual-2 | `embedding(model="voyage/voyage-multilingual-2 ", input)` | +| voyage-large-2-instruct | `embedding(model="voyage/voyage-large-2-instruct", input)` | +| voyage-large-2 | `embedding(model="voyage/voyage-large-2", input)` | +| voyage-2 | `embedding(model="voyage/voyage-2", input)` | | voyage-lite-02-instruct | `embedding(model="voyage/voyage-lite-02-instruct", input)` | -| voyage-01 | `embedding(model="voyage/voyage-01", input)` | -| voyage-lite-01 | `embedding(model="voyage/voyage-lite-01", input)` | -| voyage-lite-01-instruct | `embedding(model="voyage/voyage-lite-01-instruct", input)` | \ No newline at end of file +| voyage-01 | `embedding(model="voyage/voyage-01", input)` | +| voyage-lite-01 | `embedding(model="voyage/voyage-lite-01", input)` | +| voyage-lite-01-instruct | `embedding(model="voyage/voyage-lite-01-instruct", input)` | From 818db5e59032f37bd8f8440fd064607ae3f4c924 Mon Sep 17 00:00:00 2001 From: fzowl Date: Wed, 21 May 2025 13:09:42 +0200 Subject: [PATCH 05/21] Updating the available VoyageAI models in the docs --- docs/my-website/docs/providers/voyage.md | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/docs/my-website/docs/providers/voyage.md b/docs/my-website/docs/providers/voyage.md index 6ab6b1846f5..4b729bc9f58 100644 --- a/docs/my-website/docs/providers/voyage.md +++ b/docs/my-website/docs/providers/voyage.md @@ -25,6 +25,8 @@ All models listed here https://docs.voyageai.com/embeddings/#models-and-specific | Model Name | Function Call | |-------------------------|------------------------------------------------------------| +| voyage-3.5 | `embedding(model="voyage/voyage-3.5", input)` | +| voyage-3.5-lite | `embedding(model="voyage/voyage-3.5-lite", input)` | | voyage-3-large | `embedding(model="voyage/voyage-3-large", input)` | | voyage-3 | `embedding(model="voyage/voyage-3", input)` | | voyage-3-lite | `embedding(model="voyage/voyage-3-lite", input)` | @@ -35,8 +37,8 @@ All models listed here https://docs.voyageai.com/embeddings/#models-and-specific | voyage-multilingual-2 | `embedding(model="voyage/voyage-multilingual-2 ", input)` | | voyage-large-2-instruct | `embedding(model="voyage/voyage-large-2-instruct", input)` | | voyage-large-2 | `embedding(model="voyage/voyage-large-2", input)` | -| voyage-2 | `embedding(model="voyage/voyage-2", input)` | +| voyage-2 | `embedding(model="voyage/voyage-2", input)` | | voyage-lite-02-instruct | `embedding(model="voyage/voyage-lite-02-instruct", input)` | -| voyage-01 | `embedding(model="voyage/voyage-01", input)` | -| voyage-lite-01 | `embedding(model="voyage/voyage-lite-01", input)` | +| voyage-01 | `embedding(model="voyage/voyage-01", input)` | +| voyage-lite-01 | `embedding(model="voyage/voyage-lite-01", input)` | | voyage-lite-01-instruct | `embedding(model="voyage/voyage-lite-01-instruct", input)` | From f88dd65590460df01c49a6a58ab95ee7f28d8b58 Mon Sep 17 00:00:00 2001 From: fzowl Date: Wed, 29 Jul 2026 14:31:25 +0200 Subject: [PATCH 06/21] feat(voyage): add voyage-4 family + voyage-context-4, fix contextual list[str] input - Add voyage-4, voyage-4-large, voyage-4-lite, voyage-4-nano (/bin/bash) and voyage-context-4 to the cost/context map (both root and backup json). - Contextual embeddings: normalize flat list[str] input. Voyage rejects a flat list[str] unless input_type='query', so wrap each string as its own single-chunk document otherwise; single str -> [[str]]; list[list[str]] passes through. Keeps list[str] when input_type='query'. - Add tests for input normalization and voyage-4 family pricing. --- .../embedding/transformation_contextual.py | 33 ++++++- ...odel_prices_and_context_window_backup.json | 40 +++++++++ model_prices_and_context_window.json | 40 +++++++++ tests/llm_translation/test_voyage_ai.py | 85 +++++++++++++++++++ 4 files changed, 197 insertions(+), 1 deletion(-) diff --git a/litellm/llms/voyage/embedding/transformation_contextual.py b/litellm/llms/voyage/embedding/transformation_contextual.py index 4df2fa4ba31..2d661610eaa 100644 --- a/litellm/llms/voyage/embedding/transformation_contextual.py +++ b/litellm/llms/voyage/embedding/transformation_contextual.py @@ -106,11 +106,42 @@ class VoyageContextualEmbeddingConfig(BaseEmbeddingConfig): headers: dict, ) -> dict: return { - "inputs": input, + "inputs": self._transform_contextual_inputs(input, optional_params), "model": model, **optional_params, } + @staticmethod + def _transform_contextual_inputs( + input: Union[AllEmbeddingInputValues, List[List[str]]], + optional_params: dict, + ) -> List[List[str]]: + """ + Voyage's contextualized embeddings API expects ``inputs`` to be a + ``list[list[str]]`` (each inner list is a document made of chunks that + share context). + + It also accepts a flat ``list[str]`` - but only when + ``input_type == "query"``. In every other case a flat list must be + wrapped so each string becomes its own single-chunk document, otherwise + the API rejects the request with a 400. + + Reference: https://docs.voyageai.com/reference/contextualized-embeddings-api + """ + # Single string -> one document with a single chunk. + if isinstance(input, str): + return [[input]] + + # Flat list[str]: keep as list[str] when the API allows it + # (input_type="query"), otherwise wrap each string as its own document. + if isinstance(input, list) and all(isinstance(i, str) for i in input): + if optional_params.get("input_type") == "query": + return input # type: ignore[return-value] + return [[i] for i in input] + + # Already list[list[str]] (or another shape) -> pass through unchanged. + return input # type: ignore[return-value] + def transform_embedding_response( self, model: str, diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index c584deb683a..fc710497592 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -27345,6 +27345,38 @@ "mode": "embedding", "output_cost_per_token": 0.0 }, + "voyage/voyage-4": { + "input_cost_per_token": 6e-08, + "litellm_provider": "voyage", + "max_input_tokens": 32000, + "max_tokens": 32000, + "mode": "embedding", + "output_cost_per_token": 0.0 + }, + "voyage/voyage-4-large": { + "input_cost_per_token": 1.2e-07, + "litellm_provider": "voyage", + "max_input_tokens": 32000, + "max_tokens": 32000, + "mode": "embedding", + "output_cost_per_token": 0.0 + }, + "voyage/voyage-4-lite": { + "input_cost_per_token": 2e-08, + "litellm_provider": "voyage", + "max_input_tokens": 32000, + "max_tokens": 32000, + "mode": "embedding", + "output_cost_per_token": 0.0 + }, + "voyage/voyage-4-nano": { + "input_cost_per_token": 0.0, + "litellm_provider": "voyage", + "max_input_tokens": 32000, + "max_tokens": 32000, + "mode": "embedding", + "output_cost_per_token": 0.0 + }, "voyage/voyage-code-2": { "input_cost_per_token": 1.2e-07, "litellm_provider": "voyage", @@ -27369,6 +27401,14 @@ "mode": "embedding", "output_cost_per_token": 0.0 }, + "voyage/voyage-context-4": { + "input_cost_per_token": 1.2e-07, + "litellm_provider": "voyage", + "max_input_tokens": 120000, + "max_tokens": 120000, + "mode": "embedding", + "output_cost_per_token": 0.0 + }, "voyage/voyage-finance-2": { "input_cost_per_token": 1.2e-07, "litellm_provider": "voyage", diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index c584deb683a..fc710497592 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -27345,6 +27345,38 @@ "mode": "embedding", "output_cost_per_token": 0.0 }, + "voyage/voyage-4": { + "input_cost_per_token": 6e-08, + "litellm_provider": "voyage", + "max_input_tokens": 32000, + "max_tokens": 32000, + "mode": "embedding", + "output_cost_per_token": 0.0 + }, + "voyage/voyage-4-large": { + "input_cost_per_token": 1.2e-07, + "litellm_provider": "voyage", + "max_input_tokens": 32000, + "max_tokens": 32000, + "mode": "embedding", + "output_cost_per_token": 0.0 + }, + "voyage/voyage-4-lite": { + "input_cost_per_token": 2e-08, + "litellm_provider": "voyage", + "max_input_tokens": 32000, + "max_tokens": 32000, + "mode": "embedding", + "output_cost_per_token": 0.0 + }, + "voyage/voyage-4-nano": { + "input_cost_per_token": 0.0, + "litellm_provider": "voyage", + "max_input_tokens": 32000, + "max_tokens": 32000, + "mode": "embedding", + "output_cost_per_token": 0.0 + }, "voyage/voyage-code-2": { "input_cost_per_token": 1.2e-07, "litellm_provider": "voyage", @@ -27369,6 +27401,14 @@ "mode": "embedding", "output_cost_per_token": 0.0 }, + "voyage/voyage-context-4": { + "input_cost_per_token": 1.2e-07, + "litellm_provider": "voyage", + "max_input_tokens": 120000, + "max_tokens": 120000, + "mode": "embedding", + "output_cost_per_token": 0.0 + }, "voyage/voyage-finance-2": { "input_cost_per_token": 1.2e-07, "litellm_provider": "voyage", diff --git a/tests/llm_translation/test_voyage_ai.py b/tests/llm_translation/test_voyage_ai.py index a0b9ee0a44b..577ee808035 100644 --- a/tests/llm_translation/test_voyage_ai.py +++ b/tests/llm_translation/test_voyage_ai.py @@ -71,6 +71,26 @@ class TestVoyageAI(BaseLLMEmbeddingTest): assert response.usage.total_tokens > 0 +@pytest.mark.parametrize( + "model, expected_input_cost", + [ + ("voyage/voyage-4", 6e-08), + ("voyage/voyage-4-large", 1.2e-07), + ("voyage/voyage-4-lite", 2e-08), + ("voyage/voyage-4-nano", 0.0), + ("voyage/voyage-context-4", 1.2e-07), + ], +) +def test_voyage_4_family_pricing_registered(model, expected_input_cost): + """The voyage-4 family and voyage-context-4 must be in the cost map.""" + os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" + litellm.model_cost = litellm.get_model_cost_map(url="") + info = litellm.get_model_info(model=model) + assert info["litellm_provider"] == "voyage" + assert info["mode"] == "embedding" + assert info["input_cost_per_token"] == expected_input_cost + + def test_voyage_ai_embedding_extra_params(): """Test Voyage AI embedding with extra parameters""" try: @@ -142,6 +162,7 @@ class TestVoyageContextualEmbeddings: # Test contextual model detection assert config.is_contextualized_embeddings("voyage-context-3") is True + assert config.is_contextualized_embeddings("voyage-context-4") is True assert config.is_contextualized_embeddings("voyage-context-2") is True assert config.is_contextualized_embeddings("context-model") is True @@ -198,6 +219,70 @@ class TestVoyageContextualEmbeddings: assert transformed["model"] == "voyage-context-3" assert transformed["encoding_format"] == "float" + def test_contextual_flat_list_str_is_wrapped(self): + """Flat list[str] without input_type must be wrapped to list[list[str]].""" + from litellm.llms.voyage.embedding.transformation_contextual import ( + VoyageContextualEmbeddingConfig, + ) + + config = VoyageContextualEmbeddingConfig() + + transformed = config.transform_embedding_request( + "voyage-context-4", ["Hello", "world"], {}, {} + ) + + # Voyage rejects a flat list[str] unless input_type="query", so each + # string becomes its own single-chunk document. + assert transformed["inputs"] == [["Hello"], ["world"]] + assert transformed["model"] == "voyage-context-4" + + def test_contextual_flat_list_str_query_stays_flat(self): + """Flat list[str] with input_type='query' is sent as-is (list[str]).""" + from litellm.llms.voyage.embedding.transformation_contextual import ( + VoyageContextualEmbeddingConfig, + ) + + config = VoyageContextualEmbeddingConfig() + + transformed = config.transform_embedding_request( + "voyage-context-4", + ["Hello", "world"], + {"input_type": "query"}, + {}, + ) + + assert transformed["inputs"] == ["Hello", "world"] + assert transformed["input_type"] == "query" + + def test_contextual_single_string_is_wrapped(self): + """A single string is wrapped to a single document with a single chunk.""" + from litellm.llms.voyage.embedding.transformation_contextual import ( + VoyageContextualEmbeddingConfig, + ) + + config = VoyageContextualEmbeddingConfig() + + transformed = config.transform_embedding_request( + "voyage-context-4", "Hello", {}, {} + ) + + assert transformed["inputs"] == [["Hello"]] + + def test_contextual_nested_input_passthrough(self): + """Already-nested list[list[str]] input is passed through unchanged.""" + from litellm.llms.voyage.embedding.transformation_contextual import ( + VoyageContextualEmbeddingConfig, + ) + + config = VoyageContextualEmbeddingConfig() + + nested = [["Hello", "world"], ["Test"]] + transformed = config.transform_embedding_request( + "voyage-context-4", nested, {}, {} + ) + + assert transformed["inputs"] == nested + def test_contextual_embedding_response_transformation(self): """Test response transformation for contextual embeddings""" from litellm.llms.voyage.embedding.transformation_contextual import ( From 328b54d0042d213aee3e1e2c025480eae5a214d8 Mon Sep 17 00:00:00 2001 From: fzowl Date: Wed, 29 Jul 2026 15:22:31 +0200 Subject: [PATCH 07/21] fix(voyage): auto-chunk flat list[str] for contextual embeddings; fix mypy Rework per review: - Contextual inputs now prefer a flat list[str] and let the API auto-chunk (enable_auto_chunking=True, chunk_size=32000, input_type=document) instead of pre-wrapping into list[list[str]]. Matches the live API contract: * flat list[str] + input_type=query -> kept flat * flat list[str] otherwise / single str -> flat + auto-chunking * list[list[str]] -> passthrough - Fix mypy list-item error from the previous wrapping approach. - Update tests for the new normalization; all pass. --- .../embedding/transformation_contextual.py | 73 ++++++++++++++----- tests/llm_translation/test_voyage_ai.py | 47 ++++++++++-- 2 files changed, 94 insertions(+), 26 deletions(-) diff --git a/litellm/llms/voyage/embedding/transformation_contextual.py b/litellm/llms/voyage/embedding/transformation_contextual.py index 2d661610eaa..d555a1a2266 100644 --- a/litellm/llms/voyage/embedding/transformation_contextual.py +++ b/litellm/llms/voyage/embedding/transformation_contextual.py @@ -2,7 +2,7 @@ This module is used to transform the request and response for the Voyage contextualized embeddings API. This would be used for all the contextualized embeddings models in Voyage. """ -from typing import List, Optional, Union +from typing import Any, Dict, List, Optional, Tuple, Union import httpx @@ -98,6 +98,11 @@ class VoyageContextualEmbeddingConfig(BaseEmbeddingConfig): "Authorization": f"Bearer {api_key}", } + # Chunk size (in tokens) used when the API auto-chunks a flat ``list[str]``. + # Matches the voyage-context-4 context window so each string stays a single + # chunk instead of being split. + AUTO_CHUNK_SIZE = 32000 + def transform_embedding_request( self, model: str, @@ -105,42 +110,74 @@ class VoyageContextualEmbeddingConfig(BaseEmbeddingConfig): optional_params: dict, headers: dict, ) -> dict: + inputs, extra_params = self._transform_contextual_inputs(input, optional_params) return { - "inputs": self._transform_contextual_inputs(input, optional_params), + "inputs": inputs, "model": model, **optional_params, + **extra_params, } - @staticmethod + @classmethod def _transform_contextual_inputs( + cls, input: Union[AllEmbeddingInputValues, List[List[str]]], optional_params: dict, - ) -> List[List[str]]: + ) -> Tuple[Union[List[str], List[List[str]]], dict]: """ - Voyage's contextualized embeddings API expects ``inputs`` to be a - ``list[list[str]]`` (each inner list is a document made of chunks that - share context). + Normalize ``input`` for Voyage's contextualized embeddings API and + return ``(inputs, extra_params)`` where ``extra_params`` carries any + request fields (e.g. auto-chunking) needed for the chosen shape. - It also accepts a flat ``list[str]`` - but only when - ``input_type == "query"``. In every other case a flat list must be - wrapped so each string becomes its own single-chunk document, otherwise - the API rejects the request with a 400. + The API contract (verified against the live endpoint) is: + + - A flat ``list[str]`` is only accepted with ``input_type="query"`` or + with ``enable_auto_chunking=True`` (which itself requires + ``input_type="document"``). + - A ``list[list[str]]`` (each inner list = one document's chunks) is + always accepted. + + So we prefer to send a flat ``list[str]`` and let the API auto-chunk, + instead of pre-wrapping into ``list[list[str]]``: + + - ``str`` -> ``[str]`` + ``enable_auto_chunking`` (input_type=document) + - flat ``list[str]`` + ``input_type="query"`` -> kept flat, as-is + - flat ``list[str]`` otherwise -> kept flat + ``enable_auto_chunking`` + (input_type=document) + - ``list[list[str]]`` -> passed through unchanged Reference: https://docs.voyageai.com/reference/contextualized-embeddings-api """ - # Single string -> one document with a single chunk. + # Single string -> a one-element flat list, auto-chunked. if isinstance(input, str): - return [[input]] + return [input], cls._auto_chunk_params(optional_params) - # Flat list[str]: keep as list[str] when the API allows it - # (input_type="query"), otherwise wrap each string as its own document. + # Flat list[str]. if isinstance(input, list) and all(isinstance(i, str) for i in input): if optional_params.get("input_type") == "query": - return input # type: ignore[return-value] - return [[i] for i in input] + # The API accepts a flat query list as-is. + return input, {} # type: ignore[return-value] + # Otherwise let the API auto-chunk the flat list. + return input, cls._auto_chunk_params(optional_params) # type: ignore[return-value] # Already list[list[str]] (or another shape) -> pass through unchanged. - return input # type: ignore[return-value] + return input, {} # type: ignore[return-value] + + @classmethod + def _auto_chunk_params(cls, optional_params: dict) -> dict: + """ + Params required to send a flat ``list[str]`` to the contextualized API. + + ``enable_auto_chunking=True`` requires ``input_type="document"``, so set + it unless the caller already provided an ``input_type``. + """ + params: Dict[str, Any] = { + "enable_auto_chunking": True, + "chunk_size": cls.AUTO_CHUNK_SIZE, + } + if not optional_params.get("input_type"): + params["input_type"] = "document" + return params def transform_embedding_response( self, diff --git a/tests/llm_translation/test_voyage_ai.py b/tests/llm_translation/test_voyage_ai.py index 577ee808035..b7a2c41a979 100644 --- a/tests/llm_translation/test_voyage_ai.py +++ b/tests/llm_translation/test_voyage_ai.py @@ -219,8 +219,13 @@ class TestVoyageContextualEmbeddings: assert transformed["model"] == "voyage-context-3" assert transformed["encoding_format"] == "float" - def test_contextual_flat_list_str_is_wrapped(self): - """Flat list[str] without input_type must be wrapped to list[list[str]].""" + def test_contextual_flat_list_str_is_auto_chunked(self): + """Flat list[str] without input_type stays flat and is auto-chunked. + + The API accepts a flat list[str] only with input_type="query" or with + enable_auto_chunking=True (which requires input_type="document"), so we + keep the list flat and let the API auto-chunk it. + """ from litellm.llms.voyage.embedding.transformation_contextual import ( VoyageContextualEmbeddingConfig, ) @@ -231,10 +236,11 @@ class TestVoyageContextualEmbeddings: "voyage-context-4", ["Hello", "world"], {}, {} ) - # Voyage rejects a flat list[str] unless input_type="query", so each - # string becomes its own single-chunk document. - assert transformed["inputs"] == [["Hello"], ["world"]] + assert transformed["inputs"] == ["Hello", "world"] assert transformed["model"] == "voyage-context-4" + assert transformed["enable_auto_chunking"] is True + assert transformed["chunk_size"] == 32000 + assert transformed["input_type"] == "document" def test_contextual_flat_list_str_query_stays_flat(self): """Flat list[str] with input_type='query' is sent as-is (list[str]).""" @@ -253,9 +259,31 @@ class TestVoyageContextualEmbeddings: assert transformed["inputs"] == ["Hello", "world"] assert transformed["input_type"] == "query" + # A query list is accepted as-is, no auto-chunking needed. + assert "enable_auto_chunking" not in transformed - def test_contextual_single_string_is_wrapped(self): - """A single string is wrapped to a single document with a single chunk.""" + def test_contextual_flat_list_str_document_input_type_preserved(self): + """Explicit input_type='document' is preserved while auto-chunking.""" + from litellm.llms.voyage.embedding.transformation_contextual import ( + VoyageContextualEmbeddingConfig, + ) + + config = VoyageContextualEmbeddingConfig() + + transformed = config.transform_embedding_request( + "voyage-context-4", + ["Hello", "world"], + {"input_type": "document"}, + {}, + ) + + assert transformed["inputs"] == ["Hello", "world"] + assert transformed["input_type"] == "document" + assert transformed["enable_auto_chunking"] is True + assert transformed["chunk_size"] == 32000 + + def test_contextual_single_string_is_auto_chunked(self): + """A single string becomes a one-element flat list and is auto-chunked.""" from litellm.llms.voyage.embedding.transformation_contextual import ( VoyageContextualEmbeddingConfig, ) @@ -266,7 +294,10 @@ class TestVoyageContextualEmbeddings: "voyage-context-4", "Hello", {}, {} ) - assert transformed["inputs"] == [["Hello"]] + assert transformed["inputs"] == ["Hello"] + assert transformed["enable_auto_chunking"] is True + assert transformed["chunk_size"] == 32000 + assert transformed["input_type"] == "document" def test_contextual_nested_input_passthrough(self): """Already-nested list[list[str]] input is passed through unchanged.""" From be6487d33c43dcf28de410fe209ede88712513b8 Mon Sep 17 00:00:00 2001 From: fzowl <160063452+fzowl@users.noreply.github.com> Date: Wed, 29 Jul 2026 17:20:27 +0200 Subject: [PATCH 08/21] Update litellm/llms/voyage/embedding/transformation_contextual.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --- litellm/llms/voyage/embedding/transformation_contextual.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/litellm/llms/voyage/embedding/transformation_contextual.py b/litellm/llms/voyage/embedding/transformation_contextual.py index 84ed856e080..c4eb6f34283 100644 --- a/litellm/llms/voyage/embedding/transformation_contextual.py +++ b/litellm/llms/voyage/embedding/transformation_contextual.py @@ -148,6 +148,8 @@ class VoyageContextualEmbeddingConfig(BaseEmbeddingConfig): """ # Single string -> a one-element flat list, auto-chunked. if isinstance(input, str): + if optional_params.get("input_type") == "query": + return [input], {} return [input], cls._auto_chunk_params(optional_params) # Flat list[str]. From e8ae29f135d35bbea296d4fa935ef1820c65bb0d Mon Sep 17 00:00:00 2001 From: fzowl Date: Wed, 29 Jul 2026 17:45:08 +0200 Subject: [PATCH 09/21] fix(voyage): ruff format + add missing query-string test coverage --- .../embedding/transformation_contextual.py | 1 + tests/llm_translation/test_voyage_ai.py | 16 ++++++++++++++++ 2 files changed, 17 insertions(+) diff --git a/litellm/llms/voyage/embedding/transformation_contextual.py b/litellm/llms/voyage/embedding/transformation_contextual.py index c4eb6f34283..4d8bcc7ff80 100644 --- a/litellm/llms/voyage/embedding/transformation_contextual.py +++ b/litellm/llms/voyage/embedding/transformation_contextual.py @@ -2,6 +2,7 @@ This module is used to transform the request and response for the Voyage contextualized embeddings API. This would be used for all the contextualized embeddings models in Voyage. """ + from typing import Any, Dict, List, Optional, Tuple, Union import httpx diff --git a/tests/llm_translation/test_voyage_ai.py b/tests/llm_translation/test_voyage_ai.py index a0d291bc02f..e6c73240a0e 100644 --- a/tests/llm_translation/test_voyage_ai.py +++ b/tests/llm_translation/test_voyage_ai.py @@ -300,6 +300,22 @@ class TestVoyageContextualEmbeddings: assert transformed["chunk_size"] == 32000 assert transformed["input_type"] == "document" + def test_contextual_single_string_query_no_auto_chunk(self): + """A single string with input_type='query' is not auto-chunked.""" + from litellm.llms.voyage.embedding.transformation_contextual import ( + VoyageContextualEmbeddingConfig, + ) + + config = VoyageContextualEmbeddingConfig() + + transformed = config.transform_embedding_request( + "voyage-context-4", "Hello", {"input_type": "query"}, {} + ) + + assert transformed["inputs"] == ["Hello"] + assert transformed["input_type"] == "query" + assert "enable_auto_chunking" not in transformed + def test_contextual_nested_input_passthrough(self): """Already-nested list[list[str]] input is passed through unchanged.""" from litellm.llms.voyage.embedding.transformation_contextual import ( From bb6db24d58bb93454a77bbcb283f1f6ad4953a1d Mon Sep 17 00:00:00 2001 From: fzowl Date: Wed, 29 Jul 2026 18:37:25 +0200 Subject: [PATCH 10/21] fix(voyage): use modern type annotations to satisfy strict ruff budget (TID251, UP006) --- .../embedding/transformation_contextual.py | 32 +++++++++---------- 1 file changed, 16 insertions(+), 16 deletions(-) diff --git a/litellm/llms/voyage/embedding/transformation_contextual.py b/litellm/llms/voyage/embedding/transformation_contextual.py index 4d8bcc7ff80..8f029c27be7 100644 --- a/litellm/llms/voyage/embedding/transformation_contextual.py +++ b/litellm/llms/voyage/embedding/transformation_contextual.py @@ -3,7 +3,9 @@ This module is used to transform the request and response for the Voyage context This would be used for all the contextualized embeddings models in Voyage. """ -from typing import Any, Dict, List, Optional, Tuple, Union +from __future__ import annotations + +from typing import Any import httpx @@ -20,7 +22,7 @@ class VoyageError(BaseLLMException): self, status_code: int, message: str, - headers: Union[dict, httpx.Headers] = {}, + headers: dict | httpx.Headers = {}, ): self.status_code = status_code self.message = message @@ -43,12 +45,12 @@ class VoyageContextualEmbeddingConfig(BaseEmbeddingConfig): def get_complete_url( self, - api_base: Optional[str], - api_key: Optional[str], + api_base: str | None, + api_key: str | None, model: str, optional_params: dict, litellm_params: dict, - stream: Optional[bool] = None, + stream: bool | None = None, ) -> str: if api_base: if not api_base.endswith("/contextualizedembeddings"): @@ -81,11 +83,11 @@ class VoyageContextualEmbeddingConfig(BaseEmbeddingConfig): self, headers: dict, model: str, - messages: List[AllMessageValues], + messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - api_key: Optional[str] = None, - api_base: Optional[str] = None, + api_key: str | None = None, + api_base: str | None = None, ) -> dict: if api_key is None: api_key = ( @@ -105,7 +107,7 @@ class VoyageContextualEmbeddingConfig(BaseEmbeddingConfig): def transform_embedding_request( self, model: str, - input: Union[AllEmbeddingInputValues, List[List[str]]], + input: AllEmbeddingInputValues | list[list[str]], optional_params: dict, headers: dict, ) -> dict: @@ -120,9 +122,9 @@ class VoyageContextualEmbeddingConfig(BaseEmbeddingConfig): @classmethod def _transform_contextual_inputs( cls, - input: Union[AllEmbeddingInputValues, List[List[str]]], + input: AllEmbeddingInputValues | list[list[str]], optional_params: dict, - ) -> Tuple[Union[List[str], List[List[str]]], dict]: + ) -> tuple[list[str] | list[list[str]], dict]: """ Normalize ``input`` for Voyage's contextualized embeddings API and return ``(inputs, extra_params)`` where ``extra_params`` carries any @@ -172,7 +174,7 @@ class VoyageContextualEmbeddingConfig(BaseEmbeddingConfig): ``enable_auto_chunking=True`` requires ``input_type="document"``, so set it unless the caller already provided an ``input_type``. """ - params: Dict[str, Any] = { + params: dict[str, Any] = { "enable_auto_chunking": True, "chunk_size": cls.AUTO_CHUNK_SIZE, } @@ -186,7 +188,7 @@ class VoyageContextualEmbeddingConfig(BaseEmbeddingConfig): raw_response: httpx.Response, model_response: EmbeddingResponse, logging_obj: LiteLLMLoggingObj, - api_key: Optional[str] = None, + api_key: str | None = None, request_data: dict = {}, optional_params: dict = {}, litellm_params: dict = {}, @@ -208,9 +210,7 @@ class VoyageContextualEmbeddingConfig(BaseEmbeddingConfig): model_response.usage = usage return model_response - def get_error_class( - self, error_message: str, status_code: int, headers: Union[dict, httpx.Headers] - ) -> BaseLLMException: + def get_error_class(self, error_message: str, status_code: int, headers: dict | httpx.Headers) -> BaseLLMException: return VoyageError(message=error_message, status_code=status_code, headers=headers) @staticmethod From eddb2225dc6dbc79238f4e6a7815c40cb7957a6d Mon Sep 17 00:00:00 2001 From: fzowl Date: Wed, 29 Jul 2026 21:29:38 +0200 Subject: [PATCH 11/21] fix(voyage): add mutable-ok pragmas and replace type: ignore with pyright: ignore (LIT001, LIT009) --- .../embedding/transformation_contextual.py | 84 ++++++++++--------- 1 file changed, 43 insertions(+), 41 deletions(-) diff --git a/litellm/llms/voyage/embedding/transformation_contextual.py b/litellm/llms/voyage/embedding/transformation_contextual.py index 8f029c27be7..04cb58bcfcc 100644 --- a/litellm/llms/voyage/embedding/transformation_contextual.py +++ b/litellm/llms/voyage/embedding/transformation_contextual.py @@ -22,11 +22,14 @@ class VoyageError(BaseLLMException): self, status_code: int, message: str, - headers: dict | httpx.Headers = {}, + headers: dict | httpx.Headers = {}, # mutable-ok: base class signature ): self.status_code = status_code self.message = message - self.request = httpx.Request(method="POST", url="https://api.voyageai.com/v1/contextualizedembeddings") + self.request = httpx.Request( + method="POST", + url="https://api.voyageai.com/v1/contextualizedembeddings", + ) self.response = httpx.Response(status_code=status_code, request=self.request) super().__init__( status_code=status_code, @@ -48,8 +51,8 @@ class VoyageContextualEmbeddingConfig(BaseEmbeddingConfig): api_base: str | None, api_key: str | None, model: str, - optional_params: dict, - litellm_params: dict, + optional_params: dict, # mutable-ok: base class signature + litellm_params: dict, # mutable-ok: base class signature stream: bool | None = None, ) -> str: if api_base: @@ -58,16 +61,16 @@ class VoyageContextualEmbeddingConfig(BaseEmbeddingConfig): return api_base return "https://api.voyageai.com/v1/contextualizedembeddings" - def get_supported_openai_params(self, model: str) -> list: + def get_supported_openai_params(self, model: str) -> list: # mutable-ok: base class signature return ["encoding_format", "dimensions"] def map_openai_params( self, - non_default_params: dict, - optional_params: dict, + non_default_params: dict, # mutable-ok: base class signature + optional_params: dict, # mutable-ok: base class signature model: str, drop_params: bool, - ) -> dict: + ) -> dict: # mutable-ok: base class signature """ Map OpenAI params to Voyage params @@ -81,14 +84,14 @@ class VoyageContextualEmbeddingConfig(BaseEmbeddingConfig): def validate_environment( self, - headers: dict, + headers: dict, # mutable-ok: base class signature model: str, - messages: list[AllMessageValues], - optional_params: dict, - litellm_params: dict, + messages: list[AllMessageValues], # mutable-ok: base class signature + optional_params: dict, # mutable-ok: base class signature + litellm_params: dict, # mutable-ok: base class signature api_key: str | None = None, api_base: str | None = None, - ) -> dict: + ) -> dict: # mutable-ok: base class signature if api_key is None: api_key = ( get_secret_str("VOYAGE_API_KEY") @@ -99,18 +102,15 @@ class VoyageContextualEmbeddingConfig(BaseEmbeddingConfig): "Authorization": f"Bearer {api_key}", } - # Chunk size (in tokens) used when the API auto-chunks a flat ``list[str]``. - # Matches the voyage-context-4 context window so each string stays a single - # chunk instead of being split. AUTO_CHUNK_SIZE = 32000 def transform_embedding_request( self, model: str, - input: AllEmbeddingInputValues | list[list[str]], - optional_params: dict, - headers: dict, - ) -> dict: + input: AllEmbeddingInputValues | list[list[str]], # mutable-ok: base class signature + optional_params: dict, # mutable-ok: base class signature + headers: dict, # mutable-ok: base class signature + ) -> dict: # mutable-ok: base class signature inputs, extra_params = self._transform_contextual_inputs(input, optional_params) return { "inputs": inputs, @@ -122,9 +122,9 @@ class VoyageContextualEmbeddingConfig(BaseEmbeddingConfig): @classmethod def _transform_contextual_inputs( cls, - input: AllEmbeddingInputValues | list[list[str]], - optional_params: dict, - ) -> tuple[list[str] | list[list[str]], dict]: + input: AllEmbeddingInputValues | list[list[str]], # mutable-ok: union with AllEmbeddingInputValues + optional_params: dict, # mutable-ok: matches public API + ) -> tuple[list[str] | list[list[str]], dict]: # mutable-ok: returned to caller who owns it """ Normalize ``input`` for Voyage's contextualized embeddings API and return ``(inputs, extra_params)`` where ``extra_params`` carries any @@ -149,32 +149,27 @@ class VoyageContextualEmbeddingConfig(BaseEmbeddingConfig): Reference: https://docs.voyageai.com/reference/contextualized-embeddings-api """ - # Single string -> a one-element flat list, auto-chunked. if isinstance(input, str): if optional_params.get("input_type") == "query": - return [input], {} - return [input], cls._auto_chunk_params(optional_params) + return [input], {} # mutable-ok: fresh list returned to caller + return [input], cls._auto_chunk_params(optional_params) # mutable-ok: fresh list returned to caller - # Flat list[str]. if isinstance(input, list) and all(isinstance(i, str) for i in input): if optional_params.get("input_type") == "query": - # The API accepts a flat query list as-is. - return input, {} # type: ignore[return-value] - # Otherwise let the API auto-chunk the flat list. - return input, cls._auto_chunk_params(optional_params) # type: ignore[return-value] + return input, {} # pyright: ignore[reportReturnType] # narrowed to list[str] by isinstance+all check + return input, cls._auto_chunk_params(optional_params) # pyright: ignore[reportReturnType] # narrowed to list[str] - # Already list[list[str]] (or another shape) -> pass through unchanged. - return input, {} # type: ignore[return-value] + return input, {} # pyright: ignore[reportReturnType] # list[list[str]] branch @classmethod - def _auto_chunk_params(cls, optional_params: dict) -> dict: + def _auto_chunk_params(cls, optional_params: dict) -> dict: # mutable-ok: matches public API """ Params required to send a flat ``list[str]`` to the contextualized API. ``enable_auto_chunking=True`` requires ``input_type="document"``, so set it unless the caller already provided an ``input_type``. """ - params: dict[str, Any] = { + params: dict[str, Any] = { # mutable-ok: building return value "enable_auto_chunking": True, "chunk_size": cls.AUTO_CHUNK_SIZE, } @@ -189,16 +184,18 @@ class VoyageContextualEmbeddingConfig(BaseEmbeddingConfig): model_response: EmbeddingResponse, logging_obj: LiteLLMLoggingObj, api_key: str | None = None, - request_data: dict = {}, - optional_params: dict = {}, - litellm_params: dict = {}, + request_data: dict = {}, # mutable-ok: base class signature + optional_params: dict = {}, # mutable-ok: base class signature + litellm_params: dict = {}, # mutable-ok: base class signature ) -> EmbeddingResponse: try: raw_response_json = raw_response.json() except Exception: - raise VoyageError(message=raw_response.text, status_code=raw_response.status_code) + raise VoyageError( + message=raw_response.text, + status_code=raw_response.status_code, + ) - # model_response.usage model_response.model = raw_response_json.get("model") model_response.data = raw_response_json.get("data") model_response.object = raw_response_json.get("object") @@ -210,7 +207,12 @@ class VoyageContextualEmbeddingConfig(BaseEmbeddingConfig): model_response.usage = usage return model_response - def get_error_class(self, error_message: str, status_code: int, headers: dict | httpx.Headers) -> BaseLLMException: + def get_error_class( + self, + error_message: str, + status_code: int, + headers: dict | httpx.Headers, # mutable-ok: base class signature + ) -> BaseLLMException: return VoyageError(message=error_message, status_code=status_code, headers=headers) @staticmethod From af634ab4265d00f3fe730623a7f28d4ddef86d42 Mon Sep 17 00:00:00 2001 From: fzowl Date: Wed, 29 Jul 2026 22:13:15 +0200 Subject: [PATCH 12/21] test(voyage): add unit tests for contextual embedding in test_litellm (CI coverage path) --- .../test_voyage_contextual_embedding.py | 222 ++++++++++++++++++ 1 file changed, 222 insertions(+) create mode 100644 tests/test_litellm/llms/voyage/test_voyage_contextual_embedding.py diff --git a/tests/test_litellm/llms/voyage/test_voyage_contextual_embedding.py b/tests/test_litellm/llms/voyage/test_voyage_contextual_embedding.py new file mode 100644 index 00000000000..865608e286f --- /dev/null +++ b/tests/test_litellm/llms/voyage/test_voyage_contextual_embedding.py @@ -0,0 +1,222 @@ +import json +from unittest.mock import MagicMock + +import pytest + + +class TestVoyageContextualEmbeddings: + def test_contextual_model_detection(self): + from litellm.llms.voyage.embedding.transformation_contextual import ( + VoyageContextualEmbeddingConfig, + ) + + assert VoyageContextualEmbeddingConfig.is_contextualized_embeddings("voyage-context-3") + assert VoyageContextualEmbeddingConfig.is_contextualized_embeddings("voyage-context-4") + assert not VoyageContextualEmbeddingConfig.is_contextualized_embeddings("voyage-3-lite") + + def test_url_generation(self): + from litellm.llms.voyage.embedding.transformation_contextual import ( + VoyageContextualEmbeddingConfig, + ) + + config = VoyageContextualEmbeddingConfig() + assert ( + config.get_complete_url(None, None, "voyage-context-4", {}, {}) + == "https://api.voyageai.com/v1/contextualizedembeddings" + ) + assert ( + config.get_complete_url("https://custom.api.com", None, "voyage-context-4", {}, {}) + == "https://custom.api.com/contextualizedembeddings" + ) + assert ( + config.get_complete_url( + "https://custom.api.com/contextualizedembeddings", + None, + "voyage-context-4", + {}, + {}, + ) + == "https://custom.api.com/contextualizedembeddings" + ) + + def test_get_supported_openai_params(self): + from litellm.llms.voyage.embedding.transformation_contextual import ( + VoyageContextualEmbeddingConfig, + ) + + config = VoyageContextualEmbeddingConfig() + assert config.get_supported_openai_params("voyage-context-4") == [ + "encoding_format", + "dimensions", + ] + + def test_map_openai_params(self): + from litellm.llms.voyage.embedding.transformation_contextual import ( + VoyageContextualEmbeddingConfig, + ) + + config = VoyageContextualEmbeddingConfig() + result = config.map_openai_params( + {"encoding_format": "float", "dimensions": 512}, {}, "voyage-context-4", False + ) + assert result["encoding_format"] == "float" + assert result["output_dimension"] == 512 + + def test_validate_environment_with_api_key(self): + from litellm.llms.voyage.embedding.transformation_contextual import ( + VoyageContextualEmbeddingConfig, + ) + + config = VoyageContextualEmbeddingConfig() + headers = config.validate_environment( + {}, "voyage-context-4", [], {}, {}, api_key="test-key" + ) + assert headers == {"Authorization": "Bearer test-key"} + + def test_validate_environment_secret_fallback(self, monkeypatch): + import litellm.llms.voyage.embedding.transformation_contextual as module + from litellm.llms.voyage.embedding.transformation_contextual import ( + VoyageContextualEmbeddingConfig, + ) + + def fake_get_secret(name): + return "secret-key" if name == "VOYAGE_API_KEY" else None + + monkeypatch.setattr(module, "get_secret_str", fake_get_secret) + config = VoyageContextualEmbeddingConfig() + headers = config.validate_environment( + {}, "voyage-context-4", [], {}, {}, api_key=None + ) + assert headers == {"Authorization": "Bearer secret-key"} + + def test_nested_list_passthrough(self): + from litellm.llms.voyage.embedding.transformation_contextual import ( + VoyageContextualEmbeddingConfig, + ) + + config = VoyageContextualEmbeddingConfig() + nested = [["Hello", "world"], ["Test"]] + transformed = config.transform_embedding_request( + "voyage-context-4", nested, {}, {} + ) + assert transformed["inputs"] == nested + assert transformed["model"] == "voyage-context-4" + assert "enable_auto_chunking" not in transformed + + def test_flat_list_str_auto_chunked(self): + from litellm.llms.voyage.embedding.transformation_contextual import ( + VoyageContextualEmbeddingConfig, + ) + + config = VoyageContextualEmbeddingConfig() + transformed = config.transform_embedding_request( + "voyage-context-4", ["Hello", "world"], {}, {} + ) + assert transformed["inputs"] == ["Hello", "world"] + assert transformed["enable_auto_chunking"] is True + assert transformed["chunk_size"] == 32000 + assert transformed["input_type"] == "document" + + def test_flat_list_str_query_no_auto_chunk(self): + from litellm.llms.voyage.embedding.transformation_contextual import ( + VoyageContextualEmbeddingConfig, + ) + + config = VoyageContextualEmbeddingConfig() + transformed = config.transform_embedding_request( + "voyage-context-4", ["Hello", "world"], {"input_type": "query"}, {} + ) + assert transformed["inputs"] == ["Hello", "world"] + assert transformed["input_type"] == "query" + assert "enable_auto_chunking" not in transformed + + def test_flat_list_str_document_preserves_input_type(self): + from litellm.llms.voyage.embedding.transformation_contextual import ( + VoyageContextualEmbeddingConfig, + ) + + config = VoyageContextualEmbeddingConfig() + transformed = config.transform_embedding_request( + "voyage-context-4", ["Hello"], {"input_type": "document"}, {} + ) + assert transformed["input_type"] == "document" + assert transformed["enable_auto_chunking"] is True + + def test_single_string_auto_chunked(self): + from litellm.llms.voyage.embedding.transformation_contextual import ( + VoyageContextualEmbeddingConfig, + ) + + config = VoyageContextualEmbeddingConfig() + transformed = config.transform_embedding_request( + "voyage-context-4", "Hello", {}, {} + ) + assert transformed["inputs"] == ["Hello"] + assert transformed["enable_auto_chunking"] is True + assert transformed["input_type"] == "document" + + def test_single_string_query_no_auto_chunk(self): + from litellm.llms.voyage.embedding.transformation_contextual import ( + VoyageContextualEmbeddingConfig, + ) + + config = VoyageContextualEmbeddingConfig() + transformed = config.transform_embedding_request( + "voyage-context-4", "Hello", {"input_type": "query"}, {} + ) + assert transformed["inputs"] == ["Hello"] + assert transformed["input_type"] == "query" + assert "enable_auto_chunking" not in transformed + + def test_response_transformation(self): + from litellm.llms.voyage.embedding.transformation_contextual import ( + VoyageContextualEmbeddingConfig, + ) + from litellm.types.utils import EmbeddingResponse + + config = VoyageContextualEmbeddingConfig() + response_payload = { + "object": "list", + "data": [{"object": "embedding", "embedding": [0.1, 0.2], "index": 0}], + "model": "voyage-context-4", + "usage": {"total_tokens": 24}, + } + raw_response = MagicMock() + raw_response.json.return_value = response_payload + raw_response.status_code = 200 + raw_response.text = json.dumps(response_payload) + + model_response = EmbeddingResponse() + transformed = config.transform_embedding_response( + "voyage-context-4", raw_response, model_response, MagicMock() + ) + assert transformed.model == "voyage-context-4" + assert transformed.object == "list" + assert transformed.data == response_payload["data"] + assert transformed.usage.prompt_tokens == 24 + assert transformed.usage.total_tokens == 24 + + def test_error_response_and_error_class(self): + from litellm.llms.voyage.embedding.transformation_contextual import ( + VoyageContextualEmbeddingConfig, + VoyageError, + ) + from litellm.types.utils import EmbeddingResponse + + config = VoyageContextualEmbeddingConfig() + raw_response = MagicMock() + raw_response.json.side_effect = ValueError("not json") + raw_response.status_code = 400 + raw_response.text = "bad request" + + with pytest.raises(VoyageError) as exc_info: + config.transform_embedding_response( + "voyage-context-4", raw_response, EmbeddingResponse(), MagicMock() + ) + assert exc_info.value.status_code == 400 + assert exc_info.value.message == "bad request" + + error = config.get_error_class("rate limited", 429, {"x-test": "1"}) + assert isinstance(error, VoyageError) + assert error.status_code == 429 + assert error.message == "rate limited" From a1be4236322d7615eece57ed6cc84867e194e866 Mon Sep 17 00:00:00 2001 From: fzowl Date: Sat, 15 Aug 2026 20:29:31 +0200 Subject: [PATCH 13/21] fix(voyage): drop banned typing.Any from contextual transform to satisfy strict lint gate --- litellm/llms/voyage/embedding/transformation_contextual.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/litellm/llms/voyage/embedding/transformation_contextual.py b/litellm/llms/voyage/embedding/transformation_contextual.py index d473450a969..0c02f00fbd5 100644 --- a/litellm/llms/voyage/embedding/transformation_contextual.py +++ b/litellm/llms/voyage/embedding/transformation_contextual.py @@ -3,7 +3,7 @@ This module is used to transform the request and response for the Voyage context This would be used for all the contextualized embeddings models in Voyage. """ -from typing import Any, Final +from typing import Final import httpx @@ -167,7 +167,7 @@ class VoyageContextualEmbeddingConfig(BaseEmbeddingConfig): ``enable_auto_chunking=True`` requires ``input_type="document"``, so set it unless the caller already provided an ``input_type``. """ - params: dict[str, Any] = { # mutable-ok: building return value + params: dict[str, object] = { # mutable-ok: building return value "enable_auto_chunking": True, "chunk_size": cls.AUTO_CHUNK_SIZE, } From 8c8cd890adcc14ed8c13d2b315e8eebc6d5c031d Mon Sep 17 00:00:00 2001 From: fzowl Date: Sat, 15 Aug 2026 20:38:40 +0200 Subject: [PATCH 14/21] fix(voyage): drop redundant isinstance list check to satisfy basedpyright budget --- litellm/llms/voyage/embedding/transformation_contextual.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/litellm/llms/voyage/embedding/transformation_contextual.py b/litellm/llms/voyage/embedding/transformation_contextual.py index 0c02f00fbd5..a9a4a01f8e3 100644 --- a/litellm/llms/voyage/embedding/transformation_contextual.py +++ b/litellm/llms/voyage/embedding/transformation_contextual.py @@ -152,9 +152,9 @@ class VoyageContextualEmbeddingConfig(BaseEmbeddingConfig): return [input], {} # mutable-ok: fresh list returned to caller return [input], cls._auto_chunk_params(optional_params) # mutable-ok: fresh list returned to caller - if isinstance(input, list) and all(isinstance(i, str) for i in input): + if all(isinstance(i, str) for i in input): if optional_params.get("input_type") == "query": - return input, {} # pyright: ignore[reportReturnType] # narrowed to list[str] by isinstance+all check + return input, {} # pyright: ignore[reportReturnType] # narrowed to list[str] by all(isinstance) check return input, cls._auto_chunk_params(optional_params) # pyright: ignore[reportReturnType] # narrowed to list[str] return input, {} # pyright: ignore[reportReturnType] # list[list[str]] branch From f4ebcef0a17c2952ab90100fdba702d863c728f9 Mon Sep 17 00:00:00 2001 From: moe-berri Date: Thu, 10 Sep 2026 12:32:09 -0700 Subject: [PATCH 15/21] fix(router): classify encrypted delegated tasks with native Responses --- .../transformation.py | 3 + .../complexity_router/README.md | 10 + .../complexity_router/complexity_router.py | 119 ++++++++++-- .../router_strategy/test_complexity_router.py | 178 ++++++++++++++++++ 4 files changed, 293 insertions(+), 17 deletions(-) diff --git a/litellm/completion_extras/litellm_responses_transformation/transformation.py b/litellm/completion_extras/litellm_responses_transformation/transformation.py index c3c18e3d009..5a6debc4af5 100644 --- a/litellm/completion_extras/litellm_responses_transformation/transformation.py +++ b/litellm/completion_extras/litellm_responses_transformation/transformation.py @@ -1201,6 +1201,9 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): # Cast to Any to match the expected union type for tools list items tools.append(cast(Any, web_search_tool)) + def transform_response_format_to_text_format(self, response_format: object) -> "ResponseText | None": + return self._transform_response_format_to_text_format(response_format) + def _transform_response_format_to_text_format(self, response_format: object) -> "ResponseText | None": """ Transform Chat Completion response_format parameter to Responses API text.format parameter. diff --git a/litellm/router_strategy/complexity_router/README.md b/litellm/router_strategy/complexity_router/README.md index d605b43e42a..79119008aa2 100644 --- a/litellm/router_strategy/complexity_router/README.md +++ b/litellm/router_strategy/complexity_router/README.md @@ -361,6 +361,16 @@ model_list: keep the classifier deployment or provider default, or set a supported value such as `none` or `low` to override that call. +When the current ask is a Responses API `agent_message` containing `encrypted_content`, LLM +classification preserves the encrypted task and uses native Responses. This also bypasses the +local scoring shortcut in `heuristic_first` and `hybrid` modes. The configured classifier must use +a native OpenAI or Azure OpenAI Responses deployment with access to the encrypted content. The +provider handles the encrypted task, and the classifier still chooses the tier dynamically + +Unsupported classifier deployments and provider decryption errors use the existing +`classifier_fallback` policy. No fixed tier is introduced for encrypted tasks. Plaintext asks and +requests carrying only historical encrypted reasoning retain the existing classifier path + Classifier calls have a one-attempt hard deadline. After a timeout, the router opens a process-local circuit for that classifier and sends every session through `classifier_fallback` for `classifier_llm_config.circuit_breaker_cooldown_seconds` (30 seconds by default). When the cooldown diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index faafcea404a..3babe5b96ed 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -25,7 +25,7 @@ from threading import Lock from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, NamedTuple, cast -from pydantic import BaseModel, create_model +from pydantic import BaseModel, TypeAdapter, create_model from litellm._logging import verbose_router_logger from litellm.constants import ( @@ -56,6 +56,7 @@ from litellm.types.llms.openai import ( AllMessageValues, ChatCompletionImageObject, ChatCompletionTextObject, + ResponsesAPIResponse, ) from litellm.types.utils import ( AUTOROUTER_CLASSIFIER_CALL_ORIGIN, @@ -364,7 +365,7 @@ def _parent_session_kwargs(request_kwargs: Mapping[str, Any] | None) -> Mapping[ return {k: kwargs[k] for k in ("litellm_session_id", "litellm_trace_id") if kwargs.get(k) is not None} -def _response_cost_or_none(response: ModelResponse) -> float | None: +def _response_cost_or_none(response: ModelResponse | ResponsesAPIResponse) -> float | None: hidden_params: Final = response._hidden_params if not isinstance(hidden_params, dict): return None @@ -486,6 +487,36 @@ def _human_text(content: object, marker_pairs: tuple[tuple[str, str], ...] = _DE return _strip_reminder_blocks(_message_text(content), marker_pairs) +def _encrypted_classifier_task( + request_kwargs: Mapping[str, object] | None, + marker_pairs: tuple[tuple[str, str], ...], +) -> dict[str, object] | None: + from litellm.litellm_core_utils.prompt_templates.factory import resolve_structured_messages + + raw_input: Final = (request_kwargs or EMPTY_MAPPING).get("input") + if not isinstance(raw_input, list) or (request_kwargs or EMPTY_MAPPING).get("messages"): + return None + items: Final = TypeAdapter(tuple[dict[str, object], ...]).validate_python(raw_input) + current: Final = next( + ( + item + for item in reversed(items) + if (messages := resolve_structured_messages(messages=None, request_kwargs={"input": [item]})) + and any(_iter_human_asks_newest_first(messages, marker_pairs)) + ), + None, + ) + if current is None or current.get("type") != "agent_message" or not isinstance(current.get("content"), list): + return None + parts: Final = TypeAdapter(tuple[dict[str, object], ...]).validate_python(current["content"]) + if not any(part.get("type") == "encrypted_content" and part.get("encrypted_content") for part in parts): + return None + return { + **current, + "content": [part for part in parts if part.get("type") in ("input_text", "encrypted_content")], + } + + def _iter_human_asks_newest_first( messages: Sequence[Mapping[str, object]], marker_pairs: tuple[tuple[str, str], ...] = _DEFAULT_REMINDER_MARKERS, @@ -1630,6 +1661,10 @@ class ComplexityRouter(CustomLogger): return self._classify_with_heuristic_v2(prompt) if self.config.classifier_type == "custom": return await self._classify_with_plugin(prompt, system_prompt, request_kwargs, raw_messages) + if self.config.classifier_type in ("heuristic_first", "hybrid") and _encrypted_classifier_task( + request_kwargs, self._reminder_markers + ): + return await self._llm_classifier_outcome(prompt, system_prompt, request_kwargs, messages) if self.config.classifier_type == "heuristic_first" and self.config.classifier_llm_config is not None: return await self._classify_heuristic_first(prompt, system_prompt, request_kwargs, messages) if self.config.classifier_type == "hybrid" and self.config.classifier_llm_config is not None: @@ -1970,8 +2005,9 @@ class ComplexityRouter(CustomLogger): > 1 ) + encrypted_task: Final = _encrypted_classifier_task(request_kwargs, self._reminder_markers) user_payload: Final = self._build_classifier_user_payload( - prompt=prompt, + prompt="The delegated task in the following agent_message." if encrypted_task is not None else prompt, system_prompt=system_prompt, prior_turns=prior_turns, messages=messages, @@ -2004,34 +2040,37 @@ class ComplexityRouter(CustomLogger): if llm_config.reasoning_effort is not None: classifier_call_params = MappingProxyType({"reasoning_effort": llm_config.reasoning_effort}) - proxy_server_request: Final = { - "body": { - "model": llm_config.model, - "messages": messages_for_call, - "response_format": response_format, - **classifier_call_params, - } - } + payload: Final = ( + self._native_classifier_payload(llm_config.model, messages_for_call, response_format, encrypted_task) + if encrypted_task is not None + else {"messages": messages_for_call, "response_format": response_format, **classifier_call_params} + ) + proxy_server_request: Final = {"body": {"model": llm_config.model, **payload}} + classify: Final = ( + self.litellm_router_instance.aresponses + if encrypted_task is not None + else self.litellm_router_instance.acompletion + ) classifier_timeout_s: Final[float] = llm_config.timeout_ms / 1000 - response: Final[ModelResponse] = await asyncio.wait_for( - self.litellm_router_instance.acompletion( + response: Final[ModelResponse | ResponsesAPIResponse] = await asyncio.wait_for( + classify( model=llm_config.model, - messages=messages_for_call, stream=False, - response_format=response_format, timeout=classifier_timeout_s, num_retries=0, disable_fallbacks=True, metadata=metadata, proxy_server_request=proxy_server_request, turn_off_message_logging=turn_off_message_logging, - **classifier_call_params, + **payload, **_parent_session_kwargs(request_kwargs), ), timeout=classifier_timeout_s, ) - content: Final = response.choices[0].message.content + content: Final = ( + response.output_text if isinstance(response, ResponsesAPIResponse) else response.choices[0].message.content + ) if not content: raise ValueError("LLM classifier returned empty content") raw_tier: Final = _LabeledTierClassification.model_validate_json(content).tier @@ -2040,6 +2079,52 @@ class ComplexityRouter(CustomLogger): raise ValueError(f"LLM classifier returned an unrecognized tier: {raw_tier!r}") return tier, _response_cost_or_none(response) + def _native_classifier_payload( + self, + model: str, + messages: list[AllMessageValues], # mutable-ok: existing transformation accepts the SDK message list + response_format: Mapping[str, object], + encrypted_task: Mapping[str, object], + ) -> Mapping[str, object]: + from litellm.completion_extras.litellm_responses_transformation.transformation import ( + LiteLLMResponsesTransformationHandler, + ) + from litellm.litellm_core_utils.get_llm_provider_logic import declared_authenticating_provider, get_llm_provider + from litellm.types.router import LiteLLM_Params + + deployments: Final = self._group_deployments(model) + if not deployments: + raise ValueError("Encrypted task classification requires a native OpenAI Responses classifier deployment") + for params in (LiteLLM_Params.model_validate(deployment.get("litellm_params")) for deployment in deployments): + if declared_authenticating_provider(params.model, params.custom_llm_provider): + raise ValueError( + "Encrypted task classification requires a native OpenAI Responses classifier deployment" + ) + _, provider, _, _ = get_llm_provider(model=params.model, litellm_params=params) + if ( + provider not in ("openai", "azure") + or params.use_chat_completions_api + or params.model.startswith("openai/chat_completions/") + ): + raise ValueError( + "Encrypted task classification requires a native OpenAI Responses classifier deployment" + ) + transformation: Final = LiteLLMResponsesTransformationHandler() + input_items, instructions = transformation.convert_chat_completion_messages_to_responses_api(messages) + llm_config: Final = self.config.classifier_llm_config + reasoning: Final = ( + {"reasoning": {"effort": llm_config.reasoning_effort}} + if llm_config is not None and llm_config.reasoning_effort is not None + else {} + ) + return { + "input": [*input_items, encrypted_task], + "instructions": instructions, + "text": transformation.transform_response_format_to_text_format(dict(response_format)), + "store": False, + **reasoning, + } + @staticmethod def _build_classifier_user_payload( prompt: str, diff --git a/tests/test_litellm/router_strategy/test_complexity_router.py b/tests/test_litellm/router_strategy/test_complexity_router.py index 51103297c58..87438ef5c69 100644 --- a/tests/test_litellm/router_strategy/test_complexity_router.py +++ b/tests/test_litellm/router_strategy/test_complexity_router.py @@ -5,6 +5,8 @@ Tests the rule-based complexity scoring and tier assignment logic. """ import asyncio +import copy +import json import logging import sys import time @@ -66,6 +68,7 @@ from litellm.types.router import ( LiteLLM_Params, TaggedPreRoutingStrategy, ) +from litellm.types.llms.openai import ResponsesAPIResponse requires_semantic_router = pytest.mark.skipif( @@ -2482,6 +2485,181 @@ class TestTierLabels: assert set(config.tier_boundaries) == {"simple_medium", "medium_complex", "complex_reasoning"} +def _encrypted_agent_task() -> dict[str, object]: + return { + "type": "agent_message", + "author": "/root", + "recipient": "/root/child", + "content": [ + {"type": "input_text", "text": "Message Type: NEW_TASK\nTask name: /root/child\nPayload:\nHello"}, + {"type": "encrypted_content", "encrypted_content": "opaque-provider-task"}, + ], + } + + +def _native_classifier_response(content: str) -> ResponsesAPIResponse: + response: Final = ResponsesAPIResponse( + id="resp_classifier", + created_at=0, + status="completed", + output=[{"type": "message", "role": "assistant", "content": [{"type": "output_text", "text": content}]}], + ) + response._hidden_params = {"response_cost": 0.0001} + return response + + +def _native_classifier_router( + output: str = '{"tier":"REASONING"}', + classifier_type: str = "llm", + deployment_model: str = "openai/gpt-6-astra", + failure: Exception | None = None, +) -> tuple[ComplexityRouter, MagicMock]: + dependency: Final = MagicMock( + aresponses=AsyncMock(return_value=_native_classifier_response(output), side_effect=failure), + acompletion=AsyncMock(return_value=_llm_response('{"tier":"SIMPLE"}')), + get_model_list=MagicMock(return_value=[{"litellm_params": {"model": deployment_model}}]), + ) + return ( + ComplexityRouter( + model_name="encrypted-router", + litellm_router_instance=dependency, + complexity_router_config={ + "tiers": {"SIMPLE": "cheap-model", "REASONING": "deep-model"}, + "classifier_type": classifier_type, + "classifier_llm_config": {"model": "classifier", "timeout_ms": 100, "reasoning_effort": "low"}, + "heuristic_first_max_tier": "SIMPLE" if classifier_type == "heuristic_first" else None, + "hybrid_boundary_margin": 0.01 if classifier_type == "hybrid" else None, + "classifier_fallback": "default_model", + "default_model": "deep-model", + "session_affinity": False, + "deployment_affinity": False, + }, + ), + dependency, + ) + + +class TestEncryptedTaskClassifier: + @pytest.mark.asyncio + @pytest.mark.parametrize("classifier_type", ["llm", "heuristic_first", "hybrid"]) + @pytest.mark.parametrize("tier,model", [("SIMPLE", "cheap-model"), ("REASONING", "deep-model")]) + async def test_encrypted_task_routes_by_native_verdict(self, classifier_type: str, tier: str, model: str): + router, dependency = _native_classifier_router(json.dumps({"tier": tier}), classifier_type) + task: Final = _encrypted_agent_task() + request: Final = { + "input": [ + {"role": "user", "content": "Prior task context"}, + task, + {"type": "function_call_output", "call_id": "call_1", "output": "Tool output"}, + {"role": "user", "content": "Injected reminder"}, + ], + "instructions": "Caller constraints", + "tools": [{"type": "function", "name": "execute"}], + "previous_response_id": "resp_parent", + "litellm_session_id": "parent-session", + "litellm_trace_id": "parent-trace", + "turn_off_message_logging": True, + "litellm_metadata": {"user_api_key_hash": "caller-key-hash"}, + } + original: Final = copy.deepcopy(request) + + result: Final = await router.async_pre_routing_hook(model="encrypted-router", request_kwargs=request) + + assert result.model == model + assert result.routing_decision["tier"] == tier + assert result.routing_decision["cause"] == "llm_classifier" + assert result.routing_decision["classifier_cost"] == 0.0001 + assert result.messages is None + assert request == original + dependency.acompletion.assert_not_called() + call: Final = dependency.aresponses.call_args.kwargs + assert call["input"][-1] == task + assert "opaque-provider-task" not in json.dumps(call["input"][:-1]) + assert "Prior task context" in json.dumps(call["input"][:-1]) + assert "Caller constraints" in json.dumps(call["input"][:-1]) + assert "Caller constraints" not in call["instructions"] + assert "SIMPLE" in call["instructions"] and "REASONING" in call["instructions"] + assert call["text"]["format"]["schema"]["properties"]["tier"]["enum"] == [ + "SIMPLE", "MEDIUM", "COMPLEX", "REASONING" + ] + assert call["text"]["format"]["strict"] is True + assert call["reasoning"] == {"effort": "low"} + assert call["store"] is False + assert call["stream"] is False + assert "tools" not in call and "previous_response_id" not in call + assert "messages" not in call and "response_format" not in call + assert call["timeout"] == 0.1 and call["num_retries"] == 0 and call["disable_fallbacks"] is True + assert call["litellm_session_id"] == "parent-session" + assert call["litellm_trace_id"] == "parent-trace" + assert call["turn_off_message_logging"] is True + assert call["metadata"]["user_api_key_hash"] == "caller-key-hash" + assert call["proxy_server_request"]["body"]["input"] == call["input"] + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "items", + [ + [{"type": "reasoning", "encrypted_content": "opaque-history", "summary": []}, {"role": "user", "content": "hi"}], + [_encrypted_agent_task(), {"role": "user", "content": "hi"}], + [{**_encrypted_agent_task(), "content": [{"type": "input_text", "text": "hi"}]}], + [{"role": "user", "content": "gAAAA is plain text"}], + [{"role": "user", "content": "hi"}, {"type": "function_call_output", "call_id": "call_1", "output": "opaque-provider-task"}], + ], + ids=["historical-reasoning", "older-encrypted-task", "plaintext-agent", "ciphertext-looking-text", "tool-output"], + ) + async def test_other_asks_keep_chat_classifier(self, items: list[dict[str, object]]): + router, dependency = _native_classifier_router() + + result: Final = await router.async_pre_routing_hook(model="encrypted-router", request_kwargs={"input": items}) + + assert result.model == "cheap-model" + assert result.routing_decision["cause"] == "llm_classifier" + dependency.aresponses.assert_not_called() + dependency.acompletion.assert_awaited_once() + + @pytest.mark.asyncio + @pytest.mark.parametrize("output", ["", "not-json", '{"tier":"UNKNOWN"}']) + async def test_invalid_native_verdict_uses_existing_fallback(self, output: str): + router, dependency = _native_classifier_router(output=output) + + result: Final = await router.async_pre_routing_hook( + model="encrypted-router", request_kwargs={"input": [_encrypted_agent_task()]} + ) + + assert result.model == "deep-model" + assert result.routing_decision["cause"] == "default_model_fallback" + dependency.aresponses.assert_awaited_once() + dependency.acompletion.assert_not_called() + + @pytest.mark.asyncio + @pytest.mark.parametrize("deployment_model", ["anthropic/test-classifier", "openai/chat_completions/gpt-6-astra"]) + async def test_incompatible_classifier_does_not_flatten_encryption(self, deployment_model: str): + router, dependency = _native_classifier_router(deployment_model=deployment_model) + + result: Final = await router.async_pre_routing_hook( + model="encrypted-router", request_kwargs={"input": [_encrypted_agent_task()]} + ) + + assert result.model == "deep-model" + assert result.routing_decision["cause"] == "default_model_fallback" + dependency.aresponses.assert_not_called() + dependency.acompletion.assert_not_called() + + @pytest.mark.asyncio + @pytest.mark.parametrize("failure", [ValueError("invalid_encrypted_content"), TimeoutError("classifier timed out")]) + async def test_native_provider_failure_uses_existing_fallback(self, failure: Exception): + router, dependency = _native_classifier_router(failure=failure) + + result: Final = await router.async_pre_routing_hook( + model="encrypted-router", request_kwargs={"input": [_encrypted_agent_task()]} + ) + + assert result.model == "deep-model" + assert result.routing_decision["cause"] == "default_model_fallback" + dependency.aresponses.assert_awaited_once() + dependency.acompletion.assert_not_called() + + class TestLLMClassifier: """Test the LLM-based classifier path (aclassify) and its fallback behavior.""" From 208d554c008da91b74062413ad8d5d454f7746b5 Mon Sep 17 00:00:00 2001 From: moe-berri Date: Thu, 10 Sep 2026 13:04:36 -0700 Subject: [PATCH 16/21] fix(router): validate encrypted classifiers after deployment selection --- .../llms/base_llm/responses/transformation.py | 3 + .../llms/openai/responses/transformation.py | 3 + litellm/responses/main.py | 8 + .../complexity_router/README.md | 3 + .../complexity_router/complexity_router.py | 35 ++-- .../router_strategy/test_complexity_router.py | 151 ++++++++++++++++-- 6 files changed, 168 insertions(+), 35 deletions(-) diff --git a/litellm/llms/base_llm/responses/transformation.py b/litellm/llms/base_llm/responses/transformation.py index 1365941fe2a..14f00aaaa21 100644 --- a/litellm/llms/base_llm/responses/transformation.py +++ b/litellm/llms/base_llm/responses/transformation.py @@ -62,6 +62,9 @@ class BaseResponsesAPIConfig(ABC): """ return False + def supports_encrypted_agent_messages(self) -> bool: + return False + def sign_request( self, headers: dict, diff --git a/litellm/llms/openai/responses/transformation.py b/litellm/llms/openai/responses/transformation.py index 666e9b32011..833ae206024 100644 --- a/litellm/llms/openai/responses/transformation.py +++ b/litellm/llms/openai/responses/transformation.py @@ -110,6 +110,9 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): def supports_native_file_search(self) -> bool: return True + def supports_encrypted_agent_messages(self) -> bool: + return self.custom_llm_provider in (LlmProviders.OPENAI, LlmProviders.AZURE) + @staticmethod def _is_gpt_5_model(model: str) -> bool: """Return True only for actual OpenAI GPT-5 models. diff --git a/litellm/responses/main.py b/litellm/responses/main.py index 7a210cdd970..e88cd618a8b 100644 --- a/litellm/responses/main.py +++ b/litellm/responses/main.py @@ -1078,6 +1078,7 @@ def responses( litellm_call_id: Final[str | None] = kwargs.get("litellm_call_id", None) _is_async: Final = kwargs.pop("aresponses", False) is True skip_mcp_handler: Final = kwargs.pop("_skip_mcp_handler", False) + require_encrypted_task_support: Final = kwargs.pop("_require_encrypted_task_support", False) is True use_chat_completions_api = _pop_use_chat_completions_api_kw(kwargs) client_headers: Final = kwargs.get("headers") @@ -1186,6 +1187,13 @@ def responses( model, custom_llm_provider, deployment_model_info ) + if require_encrypted_task_support and ( + _bridges_to_chat_completions(responses_api_provider_config, use_chat_completions_api) + or responses_api_provider_config is None + or not responses_api_provider_config.supports_encrypted_agent_messages() + ): + raise ValueError("Encrypted task classification requires a compatible native Responses deployment") + local_vars.update(kwargs) # Map reasoning_effort (from litellm_params/proxy config) to reasoning when not set if reasoning is None and "reasoning_effort" in local_vars: diff --git a/litellm/router_strategy/complexity_router/README.md b/litellm/router_strategy/complexity_router/README.md index 267def7cc5b..38dfd143cbc 100644 --- a/litellm/router_strategy/complexity_router/README.md +++ b/litellm/router_strategy/complexity_router/README.md @@ -367,6 +367,9 @@ local scoring shortcut in `heuristic_first` and `hybrid` modes. The configured c a native OpenAI or Azure OpenAI Responses deployment with access to the encrypted content. The provider handles the encrypted task, and the classifier still chooses the tier dynamically +Compatibility is checked after normal deployment selection. A paused incompatible member of the +classifier group does not prevent an eligible compatible deployment from classifying the task + Unsupported classifier deployments and provider decryption errors use the existing `classifier_fallback` policy. No fixed tier is introduced for encrypted tasks. Plaintext asks and requests carrying only historical encrypted reasoning retain the existing classifier path diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index 77de4df3cd7..6df74b26f60 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -25,7 +25,7 @@ from threading import Lock from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, NamedTuple, cast -from pydantic import BaseModel, TypeAdapter, create_model +from pydantic import BaseModel, TypeAdapter, ValidationError, create_model from litellm._logging import verbose_router_logger from litellm.constants import ( @@ -504,7 +504,10 @@ def _encrypted_classifier_task( raw_input: Final = (request_kwargs or EMPTY_MAPPING).get("input") if not isinstance(raw_input, list) or (request_kwargs or EMPTY_MAPPING).get("messages"): return None - items: Final = TypeAdapter(tuple[dict[str, object], ...]).validate_python(raw_input) + try: + items: Final = TypeAdapter(tuple[dict[str, object], ...]).validate_python(raw_input) + except ValidationError: + return None current: Final = next( ( item @@ -516,7 +519,10 @@ def _encrypted_classifier_task( ) if current is None or current.get("type") != "agent_message" or not isinstance(current.get("content"), list): return None - parts: Final = TypeAdapter(tuple[dict[str, object], ...]).validate_python(current["content"]) + try: + parts: Final = TypeAdapter(tuple[dict[str, object], ...]).validate_python(current["content"]) + except ValidationError: + return None if not any(part.get("type") == "encrypted_content" and part.get("encrypted_content") for part in parts): return None return { @@ -2058,7 +2064,7 @@ class ComplexityRouter(CustomLogger): classifier_call_params = MappingProxyType({"reasoning_effort": llm_config.reasoning_effort}) payload: Final = ( - self._native_classifier_payload(llm_config.model, messages_for_call, response_format, encrypted_task) + self._native_classifier_payload(messages_for_call, response_format, encrypted_task) if encrypted_task is not None else {"messages": messages_for_call, "response_format": response_format, **classifier_call_params} ) @@ -2098,7 +2104,6 @@ class ComplexityRouter(CustomLogger): def _native_classifier_payload( self, - model: str, messages: list[AllMessageValues], # mutable-ok: existing transformation accepts the SDK message list response_format: Mapping[str, object], encrypted_task: Mapping[str, object], @@ -2106,26 +2111,7 @@ class ComplexityRouter(CustomLogger): from litellm.completion_extras.litellm_responses_transformation.transformation import ( LiteLLMResponsesTransformationHandler, ) - from litellm.litellm_core_utils.get_llm_provider_logic import declared_authenticating_provider, get_llm_provider - from litellm.types.router import LiteLLM_Params - deployments: Final = self._group_deployments(model) - if not deployments: - raise ValueError("Encrypted task classification requires a native OpenAI Responses classifier deployment") - for params in (LiteLLM_Params.model_validate(deployment.get("litellm_params")) for deployment in deployments): - if declared_authenticating_provider(params.model, params.custom_llm_provider): - raise ValueError( - "Encrypted task classification requires a native OpenAI Responses classifier deployment" - ) - _, provider, _, _ = get_llm_provider(model=params.model, litellm_params=params) - if ( - provider not in ("openai", "azure") - or params.use_chat_completions_api - or params.model.startswith("openai/chat_completions/") - ): - raise ValueError( - "Encrypted task classification requires a native OpenAI Responses classifier deployment" - ) transformation: Final = LiteLLMResponsesTransformationHandler() input_items, instructions = transformation.convert_chat_completion_messages_to_responses_api(messages) llm_config: Final = self.config.classifier_llm_config @@ -2139,6 +2125,7 @@ class ComplexityRouter(CustomLogger): "instructions": instructions, "text": transformation.transform_response_format_to_text_format(dict(response_format)), "store": False, + "_require_encrypted_task_support": True, **reasoning, } diff --git a/tests/test_litellm/router_strategy/test_complexity_router.py b/tests/test_litellm/router_strategy/test_complexity_router.py index 780465dad3b..2ea8a8b53fe 100644 --- a/tests/test_litellm/router_strategy/test_complexity_router.py +++ b/tests/test_litellm/router_strategy/test_complexity_router.py @@ -5,8 +5,10 @@ Tests the rule-based complexity scoring and tier assignment logic. """ import asyncio +from collections.abc import AsyncIterator import json from copy import deepcopy +from functools import partial import logging import sys import time @@ -14,6 +16,7 @@ from typing import Dict, Final, List from unittest.mock import AsyncMock, MagicMock, patch import pytest +import httpx from pydantic import ValidationError import litellm @@ -69,6 +72,7 @@ from litellm.types.router import ( TaggedPreRoutingStrategy, ) from litellm.types.llms.openai import ResponsesAPIResponse +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler requires_semantic_router = pytest.mark.skipif( @@ -2520,11 +2524,21 @@ def _native_classifier_router( classifier_type: str = "llm", deployment_model: str = "openai/gpt-6-astra", failure: Exception | None = None, + native_router: Router | None = None, + http_handler: AsyncHTTPHandler | None = None, ) -> tuple[ComplexityRouter, MagicMock]: dependency: Final = MagicMock( - aresponses=AsyncMock(return_value=_native_classifier_response(output), side_effect=failure), + aresponses=( + native_router.factory_function(partial(litellm.aresponses, client=http_handler), call_type="aresponses") + if native_router is not None + else AsyncMock(return_value=_native_classifier_response(output), side_effect=failure) + ), acompletion=AsyncMock(return_value=_llm_response('{"tier":"SIMPLE"}')), - get_model_list=MagicMock(return_value=[{"litellm_params": {"model": deployment_model}}]), + get_model_list=( + native_router.get_model_list + if native_router is not None + else MagicMock(return_value=[{"litellm_params": {"model": deployment_model}}]) + ), ) return ( ComplexityRouter( @@ -2533,7 +2547,11 @@ def _native_classifier_router( complexity_router_config={ "tiers": {"SIMPLE": "cheap-model", "REASONING": "deep-model"}, "classifier_type": classifier_type, - "classifier_llm_config": {"model": "classifier", "timeout_ms": 100, "reasoning_effort": "low"}, + "classifier_llm_config": { + "model": "classifier", + "timeout_ms": 5000 if native_router is not None else 100, + "reasoning_effort": "low", + }, "heuristic_first_max_tier": "SIMPLE" if classifier_type == "heuristic_first" else None, "hybrid_boundary_margin": 0.01 if classifier_type == "hybrid" else None, "classifier_fallback": "default_model", @@ -2546,6 +2564,18 @@ def _native_classifier_router( ) +@pytest.fixture +async def native_classifier_http() -> AsyncIterator[tuple[AsyncHTTPHandler, MagicMock]]: + respond: Final = MagicMock( + return_value=httpx.Response(200, json=_native_classifier_response('{"tier":"REASONING"}').model_dump()) + ) + async with httpx.AsyncClient(transport=httpx.MockTransport(respond)) as client: + handler: Final = AsyncHTTPHandler() + await handler.client.aclose() + handler.client = client + yield handler, respond + + class TestEncryptedTaskClassifier: @pytest.mark.asyncio @pytest.mark.parametrize("classifier_type", ["llm", "heuristic_first", "hybrid"]) @@ -2587,11 +2617,15 @@ class TestEncryptedTaskClassifier: assert "Caller constraints" not in call["instructions"] assert "SIMPLE" in call["instructions"] and "REASONING" in call["instructions"] assert call["text"]["format"]["schema"]["properties"]["tier"]["enum"] == [ - "SIMPLE", "MEDIUM", "COMPLEX", "REASONING" + "SIMPLE", + "MEDIUM", + "COMPLEX", + "REASONING", ] assert call["text"]["format"]["strict"] is True assert call["reasoning"] == {"effort": "low"} assert call["store"] is False + assert call["_require_encrypted_task_support"] is True assert call["stream"] is False assert "tools" not in call and "previous_response_id" not in call assert "messages" not in call and "response_format" not in call @@ -2606,13 +2640,25 @@ class TestEncryptedTaskClassifier: @pytest.mark.parametrize( "items", [ - [{"type": "reasoning", "encrypted_content": "opaque-history", "summary": []}, {"role": "user", "content": "hi"}], + [ + {"type": "reasoning", "encrypted_content": "opaque-history", "summary": []}, + {"role": "user", "content": "hi"}, + ], [_encrypted_agent_task(), {"role": "user", "content": "hi"}], [{**_encrypted_agent_task(), "content": [{"type": "input_text", "text": "hi"}]}], [{"role": "user", "content": "gAAAA is plain text"}], - [{"role": "user", "content": "hi"}, {"type": "function_call_output", "call_id": "call_1", "output": "opaque-provider-task"}], + [ + {"role": "user", "content": "hi"}, + {"type": "function_call_output", "call_id": "call_1", "output": "opaque-provider-task"}, + ], + ], + ids=[ + "historical-reasoning", + "older-encrypted-task", + "plaintext-agent", + "ciphertext-looking-text", + "tool-output", ], - ids=["historical-reasoning", "older-encrypted-task", "plaintext-agent", "ciphertext-looking-text", "tool-output"], ) async def test_other_asks_keep_chat_classifier(self, items: list[dict[str, object]]): router, dependency = _native_classifier_router() @@ -2639,9 +2685,28 @@ class TestEncryptedTaskClassifier: dependency.acompletion.assert_not_called() @pytest.mark.asyncio - @pytest.mark.parametrize("deployment_model", ["anthropic/test-classifier", "openai/chat_completions/gpt-6-astra"]) - async def test_incompatible_classifier_does_not_flatten_encryption(self, deployment_model: str): - router, dependency = _native_classifier_router(deployment_model=deployment_model) + @pytest.mark.parametrize( + "deployment_model", + ["anthropic/test-classifier", "openai/chat_completions/gpt-6-astra", "xai/test-classifier"], + ) + async def test_incompatible_classifier_does_not_flatten_encryption( + self, deployment_model: str, native_classifier_http: tuple[AsyncHTTPHandler, MagicMock] + ): + handler, respond = native_classifier_http + native: Final = Router( + model_list=[ + { + "model_name": "classifier", + "litellm_params": { + "model": deployment_model, + "api_key": "test-key", + "api_base": "https://classifier.test/v1", + }, + } + ], + num_retries=0, + ) + router, _ = _native_classifier_router(native_router=native, http_handler=handler) result: Final = await router.async_pre_routing_hook( model="encrypted-router", request_kwargs={"input": [_encrypted_agent_task()]} @@ -2649,8 +2714,72 @@ class TestEncryptedTaskClassifier: assert result.model == "deep-model" assert result.routing_decision["cause"] == "default_model_fallback" + respond.assert_not_called() + + @pytest.mark.asyncio + @pytest.mark.parametrize("blocked", [True, False]) + async def test_native_classifier_validates_selected_deployment( + self, blocked: bool, native_classifier_http: tuple[AsyncHTTPHandler, MagicMock] + ): + handler, respond = native_classifier_http + native: Final = Router( + model_list=[ + { + "model_name": "classifier", + "litellm_params": {"model": "anthropic/test-classifier", "api_key": "test-key", "order": 0}, + "model_info": {"id": "incompatible", "blocked": blocked}, + }, + { + "model_name": "classifier", + "litellm_params": { + "model": "openai/gpt-6-astra", + "api_key": "test-key", + "order": 1, + "api_base": "https://classifier.test/v1", + }, + "model_info": {"id": "compatible"}, + }, + ], + num_retries=0, + ) + router, _ = _native_classifier_router(native_router=native, http_handler=handler) + task: Final = _encrypted_agent_task() + + result: Final = await router.async_pre_routing_hook(model="encrypted-router", request_kwargs={"input": [task]}) + + assert result.model == "deep-model" + assert result.routing_decision["cause"] == ("llm_classifier" if blocked else "default_model_fallback") + if blocked: + respond.assert_called_once() + request: Final = respond.call_args.args[0] + assert request.url.path == "/v1/responses" + body: Final = json.loads(request.content) + assert body["input"][-1] == task + assert "_require_encrypted_task_support" not in body + else: + respond.assert_not_called() + + @pytest.mark.asyncio + @pytest.mark.parametrize("classifier_type", ["llm", "heuristic_first", "hybrid"]) + @pytest.mark.parametrize( + "input_items", + [ + ["unsupported-input-item"], + [{**_encrypted_agent_task(), "content": [{"type": "input_text", "text": "hi"}, None]}], + ], + ) + async def test_encrypted_detection_does_not_reject_other_input_shapes( + self, classifier_type: str, input_items: list[object] + ): + router, dependency = _native_classifier_router(classifier_type=classifier_type) + + result: Final = await router.aclassify("hi", request_kwargs={"input": input_items}) + + assert result.cause != "default_model_fallback" + assert result.tier == ComplexityTier.SIMPLE dependency.aresponses.assert_not_called() - dependency.acompletion.assert_not_called() + if classifier_type == "llm": + dependency.acompletion.assert_awaited_once() @pytest.mark.asyncio @pytest.mark.parametrize("failure", [ValueError("invalid_encrypted_content"), TimeoutError("classifier timed out")]) From d266b76a8bcbb6cbec408eeeb8758a7af1a92644 Mon Sep 17 00:00:00 2001 From: moe-berri Date: Thu, 10 Sep 2026 13:48:21 -0700 Subject: [PATCH 17/21] fix(router): honor Codex reminders and map classifier failures --- litellm/responses/main.py | 48 ++++++++++++++++--- .../complexity_router/complexity_router.py | 4 +- .../test_responses_api_bridge_flag.py | 24 ++++++++++ .../router_strategy/test_complexity_router.py | 34 +++++++++++++ 4 files changed, 102 insertions(+), 8 deletions(-) diff --git a/litellm/responses/main.py b/litellm/responses/main.py index e88cd618a8b..a68cd02e61b 100644 --- a/litellm/responses/main.py +++ b/litellm/responses/main.py @@ -4,10 +4,11 @@ from collections.abc import Coroutine, Generator, Iterable, Mapping from contextlib import contextmanager from dataclasses import dataclass from functools import partial -from typing import TYPE_CHECKING, Any, Final, Literal, Optional, cast +from typing import TYPE_CHECKING, Any, Final, Literal, NoReturn, Optional, TypeAlias, cast import httpx from pydantic import BaseModel +from typing_extensions import assert_never import litellm from litellm._logging import verbose_logger @@ -407,6 +408,37 @@ def _bridges_to_chat_completions( return responses_api_provider_config is None or use_chat_completions_api is True +_ResponsesCompatibilityFailure: TypeAlias = Literal["encrypted_task_unsupported"] + + +def _encrypted_task_support_failure( + responses_api_provider_config: BaseResponsesAPIConfig | None, use_chat_completions_api: bool +) -> _ResponsesCompatibilityFailure | None: + if ( + responses_api_provider_config is None + or _bridges_to_chat_completions(responses_api_provider_config, use_chat_completions_api) + or not responses_api_provider_config.supports_encrypted_agent_messages() + ): + return "encrypted_task_unsupported" + return None + + +def _raise_responses_compatibility_failure( + failure: _ResponsesCompatibilityFailure, model: str, custom_llm_provider: str | None +) -> NoReturn: + match failure: + case "encrypted_task_unsupported": + raise litellm.exception_type( + model=model, + custom_llm_provider=custom_llm_provider, + original_exception=ValueError( + "Encrypted task classification requires a compatible native Responses deployment" + ), + ) + case _: + assert_never(failure) + + def _deployment_passes_through_responses(model_info: object) -> bool: """Whether ``model_info.supported_endpoints`` opts the deployment into native ``{api_base}/responses``.""" if not isinstance(model_info, dict): @@ -1187,12 +1219,16 @@ def responses( model, custom_llm_provider, deployment_model_info ) - if require_encrypted_task_support and ( - _bridges_to_chat_completions(responses_api_provider_config, use_chat_completions_api) - or responses_api_provider_config is None - or not responses_api_provider_config.supports_encrypted_agent_messages() + if ( + require_encrypted_task_support + and ( + compatibility_failure := _encrypted_task_support_failure( + responses_api_provider_config, use_chat_completions_api + ) + ) + is not None ): - raise ValueError("Encrypted task classification requires a compatible native Responses deployment") + _raise_responses_compatibility_failure(compatibility_failure, model, custom_llm_provider) local_vars.update(kwargs) # Map reasoning_effort (from litellm_params/proxy config) to reasoning when not set diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index 6df74b26f60..214872c17c9 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -1690,7 +1690,7 @@ class ComplexityRouter(CustomLogger): if self.config.classifier_type == "custom": return await self._classify_with_plugin(prompt, system_prompt, request_kwargs, raw_messages) if self.config.classifier_type in ("heuristic_first", "hybrid") and _encrypted_classifier_task( - request_kwargs, self._reminder_markers + request_kwargs, self._reminder_markers_for_request(request_kwargs or EMPTY_MAPPING) ): return await self._llm_classifier_outcome(prompt, system_prompt, request_kwargs, messages) if self.config.classifier_type == "heuristic_first" and self.config.classifier_llm_config is not None: @@ -2028,7 +2028,7 @@ class ComplexityRouter(CustomLogger): > 1 ) - encrypted_task: Final = _encrypted_classifier_task(request_kwargs, self._reminder_markers) + encrypted_task: Final = _encrypted_classifier_task(request_kwargs, marker_pairs) user_payload: Final = self._build_classifier_user_payload( prompt="The delegated task in the following agent_message." if encrypted_task is not None else prompt, system_prompt=system_prompt, diff --git a/tests/test_litellm/responses/test_responses_api_bridge_flag.py b/tests/test_litellm/responses/test_responses_api_bridge_flag.py index 57aa2a6baa2..ed44a9f4545 100644 --- a/tests/test_litellm/responses/test_responses_api_bridge_flag.py +++ b/tests/test_litellm/responses/test_responses_api_bridge_flag.py @@ -7,10 +7,14 @@ calls so routed requests do not hit a custom api_base /v1/responses endpoint. """ from importlib import import_module +from typing import Final from unittest.mock import MagicMock, patch +import httpx +import pytest import litellm +from litellm.llms.custom_httpx.http_handler import HTTPHandler from litellm.types.llms.openai import ResponseAPIUsage, ResponsesAPIResponse from litellm.types.utils import Choices, Message, ModelResponse, Usage @@ -18,6 +22,26 @@ from litellm.types.utils import Choices, Message, ModelResponse, Usage class TestUseResponsesApiBridgeFlag: """Test that bridge opt-in forces the chat completions path.""" + @pytest.mark.parametrize("model", ["openai/chat_completions/gpt-6-astra", "xai/test-classifier"]) + def test_encrypted_classifier_rejection_preserves_public_error(self, model: str) -> None: + respond: Final = MagicMock(side_effect=AssertionError("Incompatible classifier sent an upstream request")) + with httpx.Client(transport=httpx.MockTransport(respond)) as client: + with pytest.raises( + litellm.APIConnectionError, + match="Encrypted task classification requires a compatible native Responses deployment", + ) as error: + litellm.responses( + model=model, + input="Delegated task", + api_key="test-key", + api_base="https://classifier.test/v1", + client=HTTPHandler(client=client), + _require_encrypted_task_support=True, + num_retries=0, + ) + assert error.value.status_code == 500 + respond.assert_not_called() + @patch.object( import_module("litellm.responses.main").litellm_completion_transformation_handler, "response_api_handler" ) diff --git a/tests/test_litellm/router_strategy/test_complexity_router.py b/tests/test_litellm/router_strategy/test_complexity_router.py index 2ea8a8b53fe..cd8f05d1e41 100644 --- a/tests/test_litellm/router_strategy/test_complexity_router.py +++ b/tests/test_litellm/router_strategy/test_complexity_router.py @@ -2577,6 +2577,40 @@ async def native_classifier_http() -> AsyncIterator[tuple[AsyncHTTPHandler, Magi class TestEncryptedTaskClassifier: + @pytest.mark.asyncio + @pytest.mark.parametrize("classifier_type", ["llm", "heuristic_first", "hybrid"]) + @pytest.mark.parametrize("codex", [True, False]) + @pytest.mark.parametrize( + "reminder", + [ + "cwd=/repo", + "Keep answers concise", + ], + ) + async def test_encrypted_task_detection_uses_request_reminder_markers( + self, classifier_type: str, codex: bool, reminder: str + ): + router, dependency = _native_classifier_router(classifier_type=classifier_type) + task: Final = _encrypted_agent_task() + request: Final = { + "input": [task, {"role": "user", "content": reminder}], + "metadata": {"user_agent": "codex-tui" if codex else "curl/8.7.1"}, + } + original: Final = deepcopy(request) + + result: Final = await router.async_pre_routing_hook(model="encrypted-router", request_kwargs=request) + + assert request == original + assert result.model == ("deep-model" if codex else "cheap-model") + if codex: + assert result.routing_decision["cause"] == "llm_classifier" + assert result.routing_decision["tier"] == "REASONING" + dependency.aresponses.assert_awaited_once() + assert dependency.aresponses.call_args.kwargs["input"][-1] == task + dependency.acompletion.assert_not_called() + else: + dependency.aresponses.assert_not_called() + @pytest.mark.asyncio @pytest.mark.parametrize("classifier_type", ["llm", "heuristic_first", "hybrid"]) @pytest.mark.parametrize("tier,model", [("SIMPLE", "cheap-model"), ("REASONING", "deep-model")]) From d3a0b0d45b814494e64edf82cd13b0683cd4aba8 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 10 Sep 2026 13:51:45 -0700 Subject: [PATCH 18/21] fix(voyage): let caller params win for contextual auto-chunking and drop duplicate cost map entries A flat list[str] sent to voyage-context-4 is treated as independent inputs and forwarded flat with enable_auto_chunking=True, chunk_size=32000, and input_type=document unless the caller already set input_type=query. Caller-supplied params now override the defaults instead of being clobbered. The voyage-4 family and voyage-context-4 cost map entries already exist on litellm_internal_staging, and voyage-4-nano is not served by the Voyage API, so those additions and their pricing test are dropped. --- .../embedding/transformation_contextual.py | 69 ++------- ...odel_prices_and_context_window_backup.json | 40 ------ model_prices_and_context_window.json | 40 ------ tests/llm_translation/test_voyage_ai.py | 131 ------------------ .../test_voyage_contextual_embedding.py | 51 ++++++- 5 files changed, 59 insertions(+), 272 deletions(-) diff --git a/litellm/llms/voyage/embedding/transformation_contextual.py b/litellm/llms/voyage/embedding/transformation_contextual.py index a9a4a01f8e3..870de8756bb 100644 --- a/litellm/llms/voyage/embedding/transformation_contextual.py +++ b/litellm/llms/voyage/embedding/transformation_contextual.py @@ -3,6 +3,7 @@ This module is used to transform the request and response for the Voyage context This would be used for all the contextualized embeddings models in Voyage. """ +from collections.abc import Mapping from typing import Final import httpx @@ -100,7 +101,7 @@ class VoyageContextualEmbeddingConfig(BaseEmbeddingConfig): "Authorization": f"Bearer {api_key}", } - AUTO_CHUNK_SIZE = 32000 + AUTO_CHUNK_SIZE: Final = 32000 def transform_embedding_request( self, @@ -109,71 +110,27 @@ class VoyageContextualEmbeddingConfig(BaseEmbeddingConfig): optional_params: dict, headers: dict, ) -> dict: - inputs, extra_params = self._transform_contextual_inputs(input, optional_params) return { - "inputs": inputs, + "inputs": [input] if isinstance(input, str) else input, "model": model, + **self._auto_chunk_params(input, optional_params), **optional_params, - **extra_params, } @classmethod - def _transform_contextual_inputs( + def _auto_chunk_params( cls, - input: AllEmbeddingInputValues | list[list[str]], # mutable-ok: union with AllEmbeddingInputValues - optional_params: dict, # mutable-ok: matches public API - ) -> tuple[list[str] | list[list[str]], dict]: # mutable-ok: returned to caller who owns it - """ - Normalize ``input`` for Voyage's contextualized embeddings API and - return ``(inputs, extra_params)`` where ``extra_params`` carries any - request fields (e.g. auto-chunking) needed for the chosen shape. - - The API contract (verified against the live endpoint) is: - - - A flat ``list[str]`` is only accepted with ``input_type="query"`` or - with ``enable_auto_chunking=True`` (which itself requires - ``input_type="document"``). - - A ``list[list[str]]`` (each inner list = one document's chunks) is - always accepted. - - So we prefer to send a flat ``list[str]`` and let the API auto-chunk, - instead of pre-wrapping into ``list[list[str]]``: - - - ``str`` -> ``[str]`` + ``enable_auto_chunking`` (input_type=document) - - flat ``list[str]`` + ``input_type="query"`` -> kept flat, as-is - - flat ``list[str]`` otherwise -> kept flat + ``enable_auto_chunking`` - (input_type=document) - - ``list[list[str]]`` -> passed through unchanged - - Reference: https://docs.voyageai.com/reference/contextualized-embeddings-api - """ - if isinstance(input, str): - if optional_params.get("input_type") == "query": - return [input], {} # mutable-ok: fresh list returned to caller - return [input], cls._auto_chunk_params(optional_params) # mutable-ok: fresh list returned to caller - - if all(isinstance(i, str) for i in input): - if optional_params.get("input_type") == "query": - return input, {} # pyright: ignore[reportReturnType] # narrowed to list[str] by all(isinstance) check - return input, cls._auto_chunk_params(optional_params) # pyright: ignore[reportReturnType] # narrowed to list[str] - - return input, {} # pyright: ignore[reportReturnType] # list[list[str]] branch - - @classmethod - def _auto_chunk_params(cls, optional_params: dict) -> dict: # mutable-ok: matches public API - """ - Params required to send a flat ``list[str]`` to the contextualized API. - - ``enable_auto_chunking=True`` requires ``input_type="document"``, so set - it unless the caller already provided an ``input_type``. - """ - params: dict[str, object] = { # mutable-ok: building return value + input: AllEmbeddingInputValues | list[list[str]], + optional_params: Mapping[str, object], + ) -> Mapping[str, object]: + is_flat: Final = isinstance(input, str) or all(isinstance(item, str) for item in input) + if not is_flat or optional_params.get("input_type") == "query": + return {} + return { "enable_auto_chunking": True, "chunk_size": cls.AUTO_CHUNK_SIZE, + "input_type": "document", } - if not optional_params.get("input_type"): - params["input_type"] = "document" - return params def transform_embedding_response( self, diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index fb1cef7ff58..12c3c5e2be9 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -39676,38 +39676,6 @@ "mode": "embedding", "output_cost_per_token": 0.0 }, - "voyage/voyage-4": { - "input_cost_per_token": 6e-08, - "litellm_provider": "voyage", - "max_input_tokens": 32000, - "max_tokens": 32000, - "mode": "embedding", - "output_cost_per_token": 0.0 - }, - "voyage/voyage-4-large": { - "input_cost_per_token": 1.2e-07, - "litellm_provider": "voyage", - "max_input_tokens": 32000, - "max_tokens": 32000, - "mode": "embedding", - "output_cost_per_token": 0.0 - }, - "voyage/voyage-4-lite": { - "input_cost_per_token": 2e-08, - "litellm_provider": "voyage", - "max_input_tokens": 32000, - "max_tokens": 32000, - "mode": "embedding", - "output_cost_per_token": 0.0 - }, - "voyage/voyage-4-nano": { - "input_cost_per_token": 0.0, - "litellm_provider": "voyage", - "max_input_tokens": 32000, - "max_tokens": 32000, - "mode": "embedding", - "output_cost_per_token": 0.0 - }, "voyage/voyage-code-2": { "input_cost_per_token": 1.2e-07, "litellm_provider": "voyage", @@ -39732,14 +39700,6 @@ "mode": "embedding", "output_cost_per_token": 0.0 }, - "voyage/voyage-context-4": { - "input_cost_per_token": 1.2e-07, - "litellm_provider": "voyage", - "max_input_tokens": 120000, - "max_tokens": 120000, - "mode": "embedding", - "output_cost_per_token": 0.0 - }, "voyage/voyage-finance-2": { "input_cost_per_token": 1.2e-07, "litellm_provider": "voyage", diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index fb1cef7ff58..12c3c5e2be9 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -39676,38 +39676,6 @@ "mode": "embedding", "output_cost_per_token": 0.0 }, - "voyage/voyage-4": { - "input_cost_per_token": 6e-08, - "litellm_provider": "voyage", - "max_input_tokens": 32000, - "max_tokens": 32000, - "mode": "embedding", - "output_cost_per_token": 0.0 - }, - "voyage/voyage-4-large": { - "input_cost_per_token": 1.2e-07, - "litellm_provider": "voyage", - "max_input_tokens": 32000, - "max_tokens": 32000, - "mode": "embedding", - "output_cost_per_token": 0.0 - }, - "voyage/voyage-4-lite": { - "input_cost_per_token": 2e-08, - "litellm_provider": "voyage", - "max_input_tokens": 32000, - "max_tokens": 32000, - "mode": "embedding", - "output_cost_per_token": 0.0 - }, - "voyage/voyage-4-nano": { - "input_cost_per_token": 0.0, - "litellm_provider": "voyage", - "max_input_tokens": 32000, - "max_tokens": 32000, - "mode": "embedding", - "output_cost_per_token": 0.0 - }, "voyage/voyage-code-2": { "input_cost_per_token": 1.2e-07, "litellm_provider": "voyage", @@ -39732,14 +39700,6 @@ "mode": "embedding", "output_cost_per_token": 0.0 }, - "voyage/voyage-context-4": { - "input_cost_per_token": 1.2e-07, - "litellm_provider": "voyage", - "max_input_tokens": 120000, - "max_tokens": 120000, - "mode": "embedding", - "output_cost_per_token": 0.0 - }, "voyage/voyage-finance-2": { "input_cost_per_token": 1.2e-07, "litellm_provider": "voyage", diff --git a/tests/llm_translation/test_voyage_ai.py b/tests/llm_translation/test_voyage_ai.py index e6c73240a0e..438bc4f3507 100644 --- a/tests/llm_translation/test_voyage_ai.py +++ b/tests/llm_translation/test_voyage_ai.py @@ -72,26 +72,6 @@ class TestVoyageAI(BaseLLMEmbeddingTest): assert response.usage.total_tokens > 0 -@pytest.mark.parametrize( - "model, expected_input_cost", - [ - ("voyage/voyage-4", 6e-08), - ("voyage/voyage-4-large", 1.2e-07), - ("voyage/voyage-4-lite", 2e-08), - ("voyage/voyage-4-nano", 0.0), - ("voyage/voyage-context-4", 1.2e-07), - ], -) -def test_voyage_4_family_pricing_registered(model, expected_input_cost): - """The voyage-4 family and voyage-context-4 must be in the cost map.""" - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") - info = litellm.get_model_info(model=model) - assert info["litellm_provider"] == "voyage" - assert info["mode"] == "embedding" - assert info["input_cost_per_token"] == expected_input_cost - - def test_voyage_ai_embedding_extra_params(): """Test Voyage AI embedding with extra parameters""" try: @@ -220,117 +200,6 @@ class TestVoyageContextualEmbeddings: assert transformed["model"] == "voyage-context-3" assert transformed["encoding_format"] == "float" - def test_contextual_flat_list_str_is_auto_chunked(self): - """Flat list[str] without input_type stays flat and is auto-chunked. - - The API accepts a flat list[str] only with input_type="query" or with - enable_auto_chunking=True (which requires input_type="document"), so we - keep the list flat and let the API auto-chunk it. - """ - from litellm.llms.voyage.embedding.transformation_contextual import ( - VoyageContextualEmbeddingConfig, - ) - - config = VoyageContextualEmbeddingConfig() - - transformed = config.transform_embedding_request( - "voyage-context-4", ["Hello", "world"], {}, {} - ) - - assert transformed["inputs"] == ["Hello", "world"] - assert transformed["model"] == "voyage-context-4" - assert transformed["enable_auto_chunking"] is True - assert transformed["chunk_size"] == 32000 - assert transformed["input_type"] == "document" - - def test_contextual_flat_list_str_query_stays_flat(self): - """Flat list[str] with input_type='query' is sent as-is (list[str]).""" - from litellm.llms.voyage.embedding.transformation_contextual import ( - VoyageContextualEmbeddingConfig, - ) - - config = VoyageContextualEmbeddingConfig() - - transformed = config.transform_embedding_request( - "voyage-context-4", - ["Hello", "world"], - {"input_type": "query"}, - {}, - ) - - assert transformed["inputs"] == ["Hello", "world"] - assert transformed["input_type"] == "query" - # A query list is accepted as-is, no auto-chunking needed. - assert "enable_auto_chunking" not in transformed - - def test_contextual_flat_list_str_document_input_type_preserved(self): - """Explicit input_type='document' is preserved while auto-chunking.""" - from litellm.llms.voyage.embedding.transformation_contextual import ( - VoyageContextualEmbeddingConfig, - ) - - config = VoyageContextualEmbeddingConfig() - - transformed = config.transform_embedding_request( - "voyage-context-4", - ["Hello", "world"], - {"input_type": "document"}, - {}, - ) - - assert transformed["inputs"] == ["Hello", "world"] - assert transformed["input_type"] == "document" - assert transformed["enable_auto_chunking"] is True - assert transformed["chunk_size"] == 32000 - - def test_contextual_single_string_is_auto_chunked(self): - """A single string becomes a one-element flat list and is auto-chunked.""" - from litellm.llms.voyage.embedding.transformation_contextual import ( - VoyageContextualEmbeddingConfig, - ) - - config = VoyageContextualEmbeddingConfig() - - transformed = config.transform_embedding_request( - "voyage-context-4", "Hello", {}, {} - ) - - assert transformed["inputs"] == ["Hello"] - assert transformed["enable_auto_chunking"] is True - assert transformed["chunk_size"] == 32000 - assert transformed["input_type"] == "document" - - def test_contextual_single_string_query_no_auto_chunk(self): - """A single string with input_type='query' is not auto-chunked.""" - from litellm.llms.voyage.embedding.transformation_contextual import ( - VoyageContextualEmbeddingConfig, - ) - - config = VoyageContextualEmbeddingConfig() - - transformed = config.transform_embedding_request( - "voyage-context-4", "Hello", {"input_type": "query"}, {} - ) - - assert transformed["inputs"] == ["Hello"] - assert transformed["input_type"] == "query" - assert "enable_auto_chunking" not in transformed - - def test_contextual_nested_input_passthrough(self): - """Already-nested list[list[str]] input is passed through unchanged.""" - from litellm.llms.voyage.embedding.transformation_contextual import ( - VoyageContextualEmbeddingConfig, - ) - - config = VoyageContextualEmbeddingConfig() - - nested = [["Hello", "world"], ["Test"]] - transformed = config.transform_embedding_request( - "voyage-context-4", nested, {}, {} - ) - - assert transformed["inputs"] == nested - def test_contextual_embedding_response_transformation(self): """Test response transformation for contextual embeddings""" from litellm.llms.voyage.embedding.transformation_contextual import ( diff --git a/tests/test_litellm/llms/voyage/test_voyage_contextual_embedding.py b/tests/test_litellm/llms/voyage/test_voyage_contextual_embedding.py index 865608e286f..d3e912ce6af 100644 --- a/tests/test_litellm/llms/voyage/test_voyage_contextual_embedding.py +++ b/tests/test_litellm/llms/voyage/test_voyage_contextual_embedding.py @@ -74,15 +74,11 @@ class TestVoyageContextualEmbeddings: assert headers == {"Authorization": "Bearer test-key"} def test_validate_environment_secret_fallback(self, monkeypatch): - import litellm.llms.voyage.embedding.transformation_contextual as module from litellm.llms.voyage.embedding.transformation_contextual import ( VoyageContextualEmbeddingConfig, ) - def fake_get_secret(name): - return "secret-key" if name == "VOYAGE_API_KEY" else None - - monkeypatch.setattr(module, "get_secret_str", fake_get_secret) + monkeypatch.setenv("VOYAGE_API_KEY", "secret-key") config = VoyageContextualEmbeddingConfig() headers = config.validate_environment( {}, "voyage-context-4", [], {}, {}, api_key=None @@ -142,6 +138,51 @@ class TestVoyageContextualEmbeddings: assert transformed["input_type"] == "document" assert transformed["enable_auto_chunking"] is True + def test_flat_list_str_caller_chunk_params_win(self): + from litellm.llms.voyage.embedding.transformation_contextual import ( + VoyageContextualEmbeddingConfig, + ) + + config = VoyageContextualEmbeddingConfig() + transformed = config.transform_embedding_request( + "voyage-context-4", + ["Hello", "world"], + {"input_type": "document", "chunk_size": 512, "chunk_overlap": 32}, + {}, + ) + assert transformed["enable_auto_chunking"] is True + assert transformed["chunk_size"] == 512 + assert transformed["chunk_overlap"] == 32 + assert transformed["input_type"] == "document" + + def test_flat_list_str_caller_can_disable_auto_chunking(self): + from litellm.llms.voyage.embedding.transformation_contextual import ( + VoyageContextualEmbeddingConfig, + ) + + config = VoyageContextualEmbeddingConfig() + transformed = config.transform_embedding_request( + "voyage-context-4", ["Hello"], {"enable_auto_chunking": False}, {} + ) + assert transformed["enable_auto_chunking"] is False + assert transformed["input_type"] == "document" + + def test_nested_list_keeps_caller_params(self): + from litellm.llms.voyage.embedding.transformation_contextual import ( + VoyageContextualEmbeddingConfig, + ) + + config = VoyageContextualEmbeddingConfig() + transformed = config.transform_embedding_request( + "voyage-context-4", [["Hello", "world"]], {"input_type": "document", "output_dimension": 512}, {} + ) + assert transformed == { + "inputs": [["Hello", "world"]], + "model": "voyage-context-4", + "input_type": "document", + "output_dimension": 512, + } + def test_single_string_auto_chunked(self): from litellm.llms.voyage.embedding.transformation_contextual import ( VoyageContextualEmbeddingConfig, From c59fc6dc28cfe2ede38603f6c6f9935049acf1e2 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 10 Sep 2026 14:25:30 -0700 Subject: [PATCH 19/21] fix(mcp): reject initialize with 403 when the key grants no MCP servers (#40616) * fix(mcp): reject initialize with 403 when the key grants no MCP servers Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(mcp): e2e expects 403 initialize for a key with no MCP servers Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(mcp): mention IP filtering in the no-servers initialize denial and keep zero-grant tool coverage Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: mateo Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../proxy/_experimental/mcp_server/server.py | 28 +++++ tests/mcp_tests/test_proxy_mcp_e2e.py | 113 ++++++++++++------ .../mcp_server/test_mcp_server.py | 87 ++++++++++++++ .../mcp_server/test_mcp_stale_session.py | 8 +- 4 files changed, 195 insertions(+), 41 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index d259fa32cd9..95129d9aeed 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -2048,18 +2048,44 @@ if MCP_AVAILABLE: return texts[0][1] return "\n\n---\n\n".join(f"[{lbl}]\n{txt}" for lbl, txt in texts) + async def _raise_if_initialize_grants_no_mcp_servers( + allowed: Sequence[MCPServer], + user_api_key_auth: UserAPIKeyAuth | None, + mcp_servers: Sequence[str] | None, + client_ip: str | None, + ) -> None: + if allowed or user_api_key_auth is None or not user_api_key_auth.api_key: + return + if mcp_servers: + await raise_denied_scoped_mcp_access( + requested_names=mcp_servers, + user_api_key_auth=user_api_key_auth, + client_ip=client_ip, + ) + no_servers_denial: Final[_McpDeniedDetail] = { + "error": ( + "The key has no MCP servers granted, or none of its granted servers is loaded and allowed for " + "this client IP. Grant servers or access groups to the key, its team, or its organization " + "(object_permission.mcp_servers), check the server's allowed IPs, and reconnect." + ) + } + raise HTTPException(status_code=403, detail=no_servers_denial) + @contextlib.asynccontextmanager async def _gateway_initialize_instructions_request_scope( user_api_key_auth: UserAPIKeyAuth | None, mcp_servers: list[str] | None, client_ip: str | None, scoped_server_endpoint: bool = False, + is_initialize: bool = False, ) -> AsyncIterator[None]: allowed: Final = await _get_allowed_mcp_servers( user_api_key_auth=user_api_key_auth, mcp_servers=mcp_servers, client_ip=client_ip, ) + if is_initialize: + await _raise_if_initialize_grants_no_mcp_servers(allowed, user_api_key_auth, mcp_servers, client_ip) if allowed: # return_exceptions=True: a per-server probe failure (incl. CancelledError # bubbled from anyio task group teardown on connection refused) must not @@ -4683,6 +4709,7 @@ if MCP_AVAILABLE: mcp_servers, _client_ip, scoped_server_endpoint=scoped_server_endpoint, + is_initialize=is_initialize, ): await target_manager.handle_request(scope, receive, local_send) if use_stateful and session_id and scope.get("method") == "DELETE": @@ -4819,6 +4846,7 @@ if MCP_AVAILABLE: mcp_servers, _sse_client_ip, scoped_server_endpoint=scoped_server_endpoint, + is_initialize=scope.get("method") == "GET", ): await sse_session_manager.handle_request(scope, receive, send) except MCPUpstreamAuthError as e: diff --git a/tests/mcp_tests/test_proxy_mcp_e2e.py b/tests/mcp_tests/test_proxy_mcp_e2e.py index f0fc3d892e2..e1099fe0a62 100644 --- a/tests/mcp_tests/test_proxy_mcp_e2e.py +++ b/tests/mcp_tests/test_proxy_mcp_e2e.py @@ -510,6 +510,37 @@ async def _call(session: ClientSession, tool_id: str, a: int = 3, b: int = 4) -> return await session.call_tool("call_tool", arguments={"tool_id": tool_id, "arguments": {"a": a, "b": b}}) +async def _raw_rpc( + proxy_server_url: str, key: str | None, method: str, params: dict[str, object], **headers: str +) -> httpx.Response: + async with httpx.AsyncClient() as client: + return await client.post( + f"{proxy_server_url}/mcp/proxy", + headers={ + "Accept": "application/json, text/event-stream", + **({"Authorization": f"Bearer {key}"} if key else {}), + **headers, + }, + json={"jsonrpc": "2.0", "id": 1, "method": method, "params": params}, + ) + + +async def _raw_initialize(proxy_server_url: str, key: str | None) -> httpx.Response: + return await _raw_rpc( + proxy_server_url, + key, + "initialize", + {"protocolVersion": "2025-03-26", "capabilities": {}, "clientInfo": {"name": "auth-test", "version": "1"}}, + ) + + +def _rpc_result(response: httpx.Response) -> dict[str, typing.Any]: + if response.headers["content-type"].startswith("text/event-stream"): + data_line = next(line for line in response.text.splitlines() if line.startswith("data:")) + return json.loads(data_line.removeprefix("data:"))["result"] + return response.json()["result"] + + def _assert_unauthorized(result: CallToolResult) -> None: assert result.isError is True assert result.content[0].text == "Unknown or unauthorized tool_id" @@ -527,13 +558,31 @@ class TestProxyMcpAuthorizationScope: _assert_unauthorized(await _call(ungranted, restricted_id)) @pytest.mark.asyncio - async def test_no_mcp_servers_sentinel_hides_every_tool(self, proxy_server_url: str) -> None: + async def test_no_mcp_servers_sentinel_rejects_initialize_and_hides_every_tool(self, proxy_server_url: str) -> None: async with _scoped_session(proxy_server_url) as granted: tool_id = (await _search(granted, "add"))["math_stdio-add"] - async with _scoped_session(proxy_server_url, "sk-none") as session: - assert await _search(session, "add") == {} - _assert_unauthorized(await session.call_tool("get_tool_schema", {"tool_id": tool_id})) - _assert_unauthorized(await _call(session, tool_id)) + response = await _raw_initialize(proxy_server_url, "sk-none") + assert response.status_code == 403, response.text + assert "no MCP servers granted" in response.json()["detail"]["error"] + + async def raw_call(name: str, arguments: dict[str, object]) -> dict[str, typing.Any]: + call = await _raw_rpc(proxy_server_url, "sk-none", "tools/call", {"name": name, "arguments": arguments}) + assert call.status_code == 200, call.text + return _rpc_result(call) + + listed = await _raw_rpc(proxy_server_url, "sk-none", "tools/list", {}) + assert listed.status_code == 200, listed.text + assert {tool["name"] for tool in _rpc_result(listed)["tools"]} == {"search_tools", "get_tool_schema", "call_tool"} + search = await raw_call("search_tools", {"query": "add"}) + assert search["isError"] is False, search + assert json.loads(search["content"][0]["text"]) == [] + for name, arguments in ( + ("get_tool_schema", {"tool_id": tool_id}), + ("call_tool", {"tool_id": tool_id, "arguments": {"a": 3, "b": 4}}), + ): + denied = await raw_call(name, arguments) + assert denied["isError"] is True, denied + assert denied["content"][0]["text"] == "Unknown or unauthorized tool_id" @pytest.mark.asyncio async def test_tool_grant_hides_ungranted_tools_and_blocks_their_ids(self, proxy_server_url: str) -> None: @@ -581,24 +630,7 @@ class TestProxyMcpAuthorizationScope: @pytest.mark.asyncio @pytest.mark.parametrize("key", [None, "sk-invalid"]) async def test_missing_or_invalid_key_cannot_initialize(self, proxy_server_url: str, key: str | None) -> None: - async with httpx.AsyncClient() as client: - response = await client.post( - f"{proxy_server_url}/mcp/proxy", - headers={ - "Accept": "application/json, text/event-stream", - **({"Authorization": f"Bearer {key}"} if key else {}), - }, - json={ - "jsonrpc": "2.0", - "id": 1, - "method": "initialize", - "params": { - "protocolVersion": "2025-03-26", - "capabilities": {}, - "clientInfo": {"name": "auth-test", "version": "1"}, - }, - }, - ) + response = await _raw_initialize(proxy_server_url, key) assert response.status_code == 401, response.text @pytest.mark.asyncio @@ -646,25 +678,28 @@ class TestProxyMcpAuthorizationScope: @pytest.mark.asyncio async def test_proxy_scope_exception_returns_iserror_and_emits_failure_log(self, proxy_server_url: str) -> None: - async with _scoped_session( + response = await _raw_rpc( proxy_server_url, "sk-none", + "tools/call", + {"name": "call_tool", "arguments": {"tool_id": "denied-scope", "arguments": {}}}, **{"x-mcp-servers": "math_restricted", "x-litellm-call-id": "proxy-scope-denial"}, - ) as session: - result = await session.call_tool("call_tool", {"tool_id": "denied-scope", "arguments": {}}) - assert result.isError is True - assert result.content[0].text == ( - "Error: The key is not allowed to access the requested MCP servers: math_restricted" - ) - async with asyncio.timeout(10): - while True: - payload = json.loads(await asyncio.to_thread(proxy_call_recorder.failures.get, True, 5)) - if payload["id"] == "proxy-scope-denial": - break - assert payload["call_type"] == "call_mcp_tool" - assert payload["status"] == "failure" - assert payload["response_cost"] == 0 - assert "math_restricted" in payload["error_str"] + ) + assert response.status_code == 200, response.text + result = _rpc_result(response) + assert result["isError"] is True + assert result["content"][0]["text"] == ( + "Error: The key is not allowed to access the requested MCP servers: math_restricted" + ) + async with asyncio.timeout(10): + while True: + payload = json.loads(await asyncio.to_thread(proxy_call_recorder.failures.get, True, 5)) + if payload["id"] == "proxy-scope-denial": + break + assert payload["call_type"] == "call_mcp_tool" + assert payload["status"] == "failure" + assert payload["response_cost"] == 0 + assert "math_restricted" in payload["error_str"] @pytest.mark.parametrize("arguments", ["wrong", False, None, [], 0]) def test_handler_rejects_non_object_arguments( diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py index 12805e355f9..7eb4396e60d 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py @@ -2003,6 +2003,11 @@ async def test_mcp_routing_chunked_initialize_to_stateful(): patch( "litellm.proxy._experimental.mcp_server.server.set_auth_context", ), + patch( # test-quality-ok: registry is empty in unit tests; key owns one server + "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", + new_callable=AsyncMock, + return_value=[MagicMock()], + ), patch( "litellm.proxy._experimental.mcp_server.server._SESSION_MANAGERS_INITIALIZED", True, @@ -2488,6 +2493,11 @@ async def test_initialize_request_tracks_active_session_after_response_header(): new_callable=AsyncMock, return_value=(owner_auth, None, None, None, None, None), ), + patch( # test-quality-ok: registry is empty in unit tests; key owns one server + "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", + new_callable=AsyncMock, + return_value=[MagicMock()], + ), patch( "litellm.proxy._experimental.mcp_server.server._SESSION_MANAGERS_INITIALIZED", True, @@ -2600,6 +2610,11 @@ async def test_initialize_request_with_existing_session_tracks_new_session(): {"x-new-header": "new"}, ), ), + patch( # test-quality-ok: registry is empty in unit tests; key owns one server + "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", + new_callable=AsyncMock, + return_value=[MagicMock()], + ), patch( "litellm.proxy._experimental.mcp_server.server._SESSION_MANAGERS_INITIALIZED", True, @@ -5614,6 +5629,78 @@ class TestGatewayCreateInitializationOptions: assert server.create_initialization_options().server_name == "litellm-mcp-server" + @pytest.mark.asyncio + async def test_initialize_with_no_granted_servers_returns_403(self): + from fastapi import HTTPException + + from litellm.proxy._experimental.mcp_server.server import ( + _gateway_initialize_instructions_request_scope, + ) + from litellm.proxy._types import UserAPIKeyAuth + + with patch( # test-quality-ok: grant resolution is the input under test + "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", + new_callable=AsyncMock, + return_value=[], + ): + with pytest.raises(HTTPException) as exc_info: + async with _gateway_initialize_instructions_request_scope( + user_api_key_auth=UserAPIKeyAuth(api_key="sk-no-mcp"), + mcp_servers=None, + client_ip=None, + is_initialize=True, + ): + pytest.fail("initialize must not proceed when the key grants no MCP servers") + + assert exc_info.value.status_code == 403 + assert "no MCP servers granted" in exc_info.value.detail["error"] + + @pytest.mark.asyncio + async def test_initialize_with_no_granted_scoped_servers_returns_scoped_denial(self): + from fastapi import HTTPException + + from litellm.proxy._experimental.mcp_server.server import ( + _gateway_initialize_instructions_request_scope, + ) + from litellm.proxy._types import UserAPIKeyAuth + + with patch( # test-quality-ok: grant resolution is the input under test + "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", + new_callable=AsyncMock, + return_value=[], + ): + with pytest.raises(HTTPException) as exc_info: + async with _gateway_initialize_instructions_request_scope( + user_api_key_auth=UserAPIKeyAuth(api_key="sk-no-mcp"), + mcp_servers=["grafana"], + client_ip=None, + is_initialize=True, + ): + pytest.fail("scoped initialize must not proceed when nothing resolves") + + assert exc_info.value.status_code == 403 + assert "grafana" in exc_info.value.detail["error"] + + @pytest.mark.asyncio + async def test_non_initialize_request_with_no_granted_servers_is_not_rejected_here(self): + from litellm.proxy._experimental.mcp_server.server import ( + _gateway_initialize_instructions_request_scope, + _mcp_gateway_initialize_instructions, + ) + from litellm.proxy._types import UserAPIKeyAuth + + with patch( # test-quality-ok: grant resolution is the input under test + "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", + new_callable=AsyncMock, + return_value=[], + ): + async with _gateway_initialize_instructions_request_scope( + user_api_key_auth=UserAPIKeyAuth(api_key="sk-no-mcp"), + mcp_servers=None, + client_ip=None, + ): + assert _mcp_gateway_initialize_instructions.get() is None + @pytest.mark.asyncio async def test_sse_handler_scopes_server_name_from_single_server_path(self): try: diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_stale_session.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_stale_session.py index feab179570b..9420eecd222 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_stale_session.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_stale_session.py @@ -750,8 +750,7 @@ async def test_admitted_subject_missing_stored_token_challenged_with_resource_me challenge = exc_info.value.headers["www-authenticate"] assert "authorization_uri=" not in challenge assert challenge == ( - 'Bearer resource_metadata="http://localhost:8000' - '/.well-known/oauth-protected-resource/mcp/repro_oauth_server"' + 'Bearer resource_metadata="http://localhost:8000/.well-known/oauth-protected-resource/mcp/repro_oauth_server"' ) @@ -938,6 +937,11 @@ async def test_handle_streamable_http_mcp_delegated_server_surfaces_upstream_cha "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager.get_mcp_server_by_name", return_value=delegated_server, ), + patch( # test-quality-ok: registry is empty in unit tests; key owns the delegated server + "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", + new_callable=AsyncMock, + return_value=[delegated_server], + ), patch.object( session_manager_stateful, "handle_request", From 03815cf9f61fa1780cda35927b5fa9450439b1b2 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 10 Sep 2026 14:25:42 -0700 Subject: [PATCH 20/21] feat(claude-code): accept https zip archive plugin sources for skills (#40496) Co-authored-by: yassin Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/_lazy_openapi_snapshot.json | 12 +-- .../claude_code_marketplace.py | 39 ++++++-- litellm/types/proxy/claude_code_endpoints.py | 10 ++- .../test_claude_code_marketplace.py | 77 +++++++++++++++- .../_components/add_plugin_form.test.tsx | 89 ++++++++++++++++++- .../skills/_components/add_plugin_form.tsx | 71 +++++++++++---- .../components/AIHub/SkillHubTableColumns.tsx | 2 +- .../claude_code_plugins/helpers.test.ts | 60 +++++++++++++ .../components/claude_code_plugins/helpers.ts | 30 +++++-- .../claude_code_plugins/skill_detail.tsx | 2 +- .../components/claude_code_plugins/types.ts | 3 +- ui/litellm-dashboard/src/lib/http/schema.d.ts | 18 ++-- 12 files changed, 356 insertions(+), 57 deletions(-) diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 4d5cdcc8003..3f060f1f5f6 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -5384,7 +5384,7 @@ "additionalProperties": { "type": "string" }, - "description": "Git source reference", + "description": "Plugin source reference", "title": "Source", "type": "object" }, @@ -5411,7 +5411,7 @@ "type": "object" }, "RegisterPluginRequest": { - "description": "Request body for registering a plugin in the marketplace.\n\nLiteLLM acts as a registry/discovery layer. Plugins are hosted on\nGitHub/GitLab/Bitbucket and referenced by their git source.", + "description": "Request body for registering a plugin in the marketplace.\n\nLiteLLM acts as a registry/discovery layer. Plugins are hosted on\nGitHub/GitLab/Bitbucket or as a zip archive on any https host and referenced by their source.", "properties": { "author": { "anyOf": [ @@ -5509,7 +5509,7 @@ "additionalProperties": { "type": "string" }, - "description": "Git source reference. Supported formats:\n- GitHub: {'source': 'github', 'repo': 'org/repo'}\n- Git URL: {'source': 'url', 'url': 'https://github.com/org/repo.git'}\n- Git Subdir: {'source': 'git-subdir', 'url': 'https://github.com/org/repo.git', 'path': 'plugins/plugin-name'}", + "description": "Plugin source reference. Supported formats:\n- GitHub: {'source': 'github', 'repo': 'org/repo'}\n- Git URL: {'source': 'url', 'url': 'https://github.com/org/repo.git'}\n- Git Subdir: {'source': 'git-subdir', 'url': 'https://github.com/org/repo.git', 'path': 'plugins/plugin-name'}\n- Zip archive on any https host (e.g. S3): {'source': 'archive', 'url': 'https://bucket.s3.amazonaws.com/plugin.zip', 'sha256': ''}", "title": "Source", "type": "object" }, @@ -5653,7 +5653,7 @@ "additionalProperties": { "type": "string" }, - "description": "Git source reference. Supported formats:\n- GitHub: {'source': 'github', 'repo': 'org/repo'}\n- Git URL: {'source': 'url', 'url': 'https://github.com/org/repo.git'}\n- Git Subdir: {'source': 'git-subdir', 'url': 'https://github.com/org/repo.git', 'path': 'plugins/plugin-name'}", + "description": "Plugin source reference. Supported formats:\n- GitHub: {'source': 'github', 'repo': 'org/repo'}\n- Git URL: {'source': 'url', 'url': 'https://github.com/org/repo.git'}\n- Git Subdir: {'source': 'git-subdir', 'url': 'https://github.com/org/repo.git', 'path': 'plugins/plugin-name'}\n- Zip archive on any https host (e.g. S3): {'source': 'archive', 'url': 'https://bucket.s3.amazonaws.com/plugin.zip', 'sha256': ''}", "title": "Source", "type": "object" }, @@ -5788,7 +5788,7 @@ ] }, "post": { - "description": "Register a new plugin in the LiteLLM marketplace.\n\nLiteLLM acts as a registry/discovery layer. Plugins are hosted on\nGitHub/GitLab/Bitbucket. Claude Code will clone from the git source\nwhen users install.\n\nThis endpoint is create-only and never overwrites. If a plugin with\nthe same name already exists it returns 409 Conflict; use\nPUT /claude-code/plugins/{plugin_name} to update an existing plugin.\n\nRequires a proxy admin API key.\n\nParameters:\n - name: Plugin name (kebab-case)\n - source: Git source reference (github, url, or git-subdir format)\n - version: Semantic version (optional)\n - description: Plugin description (optional)\n - author: Author information (optional)\n - homepage: Plugin homepage URL (optional)\n - keywords: Search keywords (optional)\n - category: Plugin category (optional)\n\nReturns:\n Registration status (action is always \"created\") and plugin information.\n\nExample:\n ```bash\n curl -X POST http://localhost:4000/claude-code/plugins \\\n -H \"Authorization: Bearer sk-...\" \\\n -H \"Content-Type: application/json\" \\\n -d '{\n \"name\": \"my-plugin\",\n \"source\": {\"source\": \"github\", \"repo\": \"org/my-plugin\"},\n \"version\": \"1.0.0\",\n \"description\": \"My awesome plugin\"\n }'\n ```", + "description": "Register a new plugin in the LiteLLM marketplace.\n\nLiteLLM acts as a registry/discovery layer. Plugins are hosted on\nGitHub/GitLab/Bitbucket or as a zip archive on any https host (e.g. S3).\nClaude Code clones the git source or downloads the archive when users install.\n\nThis endpoint is create-only and never overwrites. If a plugin with\nthe same name already exists it returns 409 Conflict; use\nPUT /claude-code/plugins/{plugin_name} to update an existing plugin.\n\nRequires a proxy admin API key.\n\nParameters:\n - name: Plugin name (kebab-case)\n - source: Plugin source reference (github, url, git-subdir, or archive format)\n - version: Semantic version (optional)\n - description: Plugin description (optional)\n - author: Author information (optional)\n - homepage: Plugin homepage URL (optional)\n - keywords: Search keywords (optional)\n - category: Plugin category (optional)\n\nReturns:\n Registration status (action is always \"created\") and plugin information.\n\nExample:\n ```bash\n curl -X POST http://localhost:4000/claude-code/plugins \\\n -H \"Authorization: Bearer sk-...\" \\\n -H \"Content-Type: application/json\" \\\n -d '{\n \"name\": \"my-plugin\",\n \"source\": {\"source\": \"github\", \"repo\": \"org/my-plugin\"},\n \"version\": \"1.0.0\",\n \"description\": \"My awesome plugin\"\n }'\n ```", "operationId": "register_plugin_claude_code_plugins_post", "requestBody": { "content": { @@ -5923,7 +5923,7 @@ ] }, "put": { - "description": "Update an existing plugin in the LiteLLM marketplace.\n\nThe plugin is identified by its name in the path, which is the resource\nidentity and cannot be changed here. This is a full replace, not a merge:\nthe manifest is rebuilt from the request body, so any optional field left\nout is reset to its default (e.g. an omitted version is cleared, not kept).\nSend the full desired state.\n\nReturns 404 if no plugin with the given name exists; use\nPOST /claude-code/plugins to create a new plugin.\n\nRequires a proxy admin API key.\n\nParameters:\n - plugin_name: Name of the plugin to update (path parameter)\n - source: Git source reference (github, url, or git-subdir format)\n - version: Semantic version (optional)\n - description: Plugin description (optional)\n - author: Author information (optional)\n - homepage: Plugin homepage URL (optional)\n - keywords: Search keywords (optional)\n - category: Plugin category (optional)\n\nReturns:\n Update status (action is always \"updated\") and plugin information.\n\nExample:\n ```bash\n curl -X PUT http://localhost:4000/claude-code/plugins/my-plugin \\\n -H \"Authorization: Bearer sk-...\" \\\n -H \"Content-Type: application/json\" \\\n -d '{\n \"source\": {\"source\": \"github\", \"repo\": \"org/my-plugin\"},\n \"version\": \"2.0.0\",\n \"description\": \"My awesome plugin\"\n }'\n ```", + "description": "Update an existing plugin in the LiteLLM marketplace.\n\nThe plugin is identified by its name in the path, which is the resource\nidentity and cannot be changed here. This is a full replace, not a merge:\nthe manifest is rebuilt from the request body, so any optional field left\nout is reset to its default (e.g. an omitted version is cleared, not kept).\nSend the full desired state.\n\nReturns 404 if no plugin with the given name exists; use\nPOST /claude-code/plugins to create a new plugin.\n\nRequires a proxy admin API key.\n\nParameters:\n - plugin_name: Name of the plugin to update (path parameter)\n - source: Plugin source reference (github, url, git-subdir, or archive format)\n - version: Semantic version (optional)\n - description: Plugin description (optional)\n - author: Author information (optional)\n - homepage: Plugin homepage URL (optional)\n - keywords: Search keywords (optional)\n - category: Plugin category (optional)\n\nReturns:\n Update status (action is always \"updated\") and plugin information.\n\nExample:\n ```bash\n curl -X PUT http://localhost:4000/claude-code/plugins/my-plugin \\\n -H \"Authorization: Bearer sk-...\" \\\n -H \"Content-Type: application/json\" \\\n -d '{\n \"source\": {\"source\": \"github\", \"repo\": \"org/my-plugin\"},\n \"version\": \"2.0.0\",\n \"description\": \"My awesome plugin\"\n }'\n ```", "operationId": "update_plugin_claude_code_plugins__plugin_name__put", "parameters": [ { diff --git a/litellm/proxy/anthropic_endpoints/claude_code_endpoints/claude_code_marketplace.py b/litellm/proxy/anthropic_endpoints/claude_code_endpoints/claude_code_marketplace.py index 65bc46edfaf..cdeb1c6a111 100644 --- a/litellm/proxy/anthropic_endpoints/claude_code_endpoints/claude_code_marketplace.py +++ b/litellm/proxy/anthropic_endpoints/claude_code_endpoints/claude_code_marketplace.py @@ -2,8 +2,9 @@ CLAUDE CODE MARKETPLACE Provides a registry/discovery layer for Claude Code plugins. -Plugins are stored as metadata + git source references in LiteLLM database. -Actual plugin files are hosted on GitHub/GitLab/Bitbucket. +Plugins are stored as metadata + source references in LiteLLM database. +Actual plugin files are hosted on GitHub/GitLab/Bitbucket or as a zip archive on +any HTTPS host (S3, Artifactory, a static file server). Endpoints: /claude-code/marketplace.json - GET - List plugins for Claude Code discovery (unauthenticated) @@ -21,6 +22,7 @@ import re from collections.abc import Mapping, Sequence from datetime import datetime, timezone from typing import Annotated, Final, Protocol, TypedDict +from urllib.parse import urlsplit from fastapi import APIRouter, Depends, HTTPException from fastapi.responses import JSONResponse @@ -162,6 +164,15 @@ async def get_marketplace(): # alphanumeric characters, dots, hyphens, and underscores. # This implicitly blocks '..', leading '/', backslashes, and percent-encoded sequences. _VALID_GIT_SUBDIR_PATH_RE: Final = re.compile(r"^[a-zA-Z0-9][a-zA-Z0-9._-]*(/[a-zA-Z0-9][a-zA-Z0-9._-]*)*$") +_VALID_SHA256_RE: Final = re.compile(r"^[0-9a-fA-F]{64}$") + + +def _is_https_url_with_host(url: str) -> bool: + try: + parts: Final = urlsplit(url) + except ValueError: + return False + return parts.scheme == "https" and bool(parts.hostname) def _validate_plugin_source(source: Mapping[str, str]) -> None: @@ -199,10 +210,24 @@ def _validate_plugin_source(source: Mapping[str, str]) -> None: "error": "git-subdir 'path' must be a relative path of the form 'segment/segment' (alphanumeric, dots, hyphens, underscores only)" }, ) + elif source_type == "archive": + if not _is_https_url_with_host(source.get("url", "")): + raise HTTPException( + status_code=400, + detail={ + "error": "archive source must include an https 'url' field " + "(e.g., 'https://bucket.s3.amazonaws.com/plugins/plugin-name.zip')" + }, + ) + if "sha256" in source and not _VALID_SHA256_RE.match(source["sha256"]): + raise HTTPException( + status_code=400, + detail={"error": "archive 'sha256' must be a 64-character hex digest"}, + ) else: raise HTTPException( status_code=400, - detail={"error": "source.source must be 'github', 'url', or 'git-subdir'"}, + detail={"error": "source.source must be 'github', 'url', 'git-subdir', or 'archive'"}, ) @@ -248,8 +273,8 @@ async def register_plugin( Register a new plugin in the LiteLLM marketplace. LiteLLM acts as a registry/discovery layer. Plugins are hosted on - GitHub/GitLab/Bitbucket. Claude Code will clone from the git source - when users install. + GitHub/GitLab/Bitbucket or as a zip archive on any https host (e.g. S3). + Claude Code clones the git source or downloads the archive when users install. This endpoint is create-only and never overwrites. If a plugin with the same name already exists it returns 409 Conflict; use @@ -259,7 +284,7 @@ async def register_plugin( Parameters: - name: Plugin name (kebab-case) - - source: Git source reference (github, url, or git-subdir format) + - source: Plugin source reference (github, url, git-subdir, or archive format) - version: Semantic version (optional) - description: Plugin description (optional) - author: Author information (optional) @@ -503,7 +528,7 @@ async def update_plugin( Parameters: - plugin_name: Name of the plugin to update (path parameter) - - source: Git source reference (github, url, or git-subdir format) + - source: Plugin source reference (github, url, git-subdir, or archive format) - version: Semantic version (optional) - description: Plugin description (optional) - author: Author information (optional) diff --git a/litellm/types/proxy/claude_code_endpoints.py b/litellm/types/proxy/claude_code_endpoints.py index 2ee1bbbbb98..dcb5561cfeb 100644 --- a/litellm/types/proxy/claude_code_endpoints.py +++ b/litellm/types/proxy/claude_code_endpoints.py @@ -25,10 +25,12 @@ class PluginSpec(BaseModel): source: dict[str, str] = Field( ..., description=( - "Git source reference. Supported formats:\n" + "Plugin source reference. Supported formats:\n" "- GitHub: {'source': 'github', 'repo': 'org/repo'}\n" "- Git URL: {'source': 'url', 'url': 'https://github.com/org/repo.git'}\n" - "- Git Subdir: {'source': 'git-subdir', 'url': 'https://github.com/org/repo.git', 'path': 'plugins/plugin-name'}" + "- Git Subdir: {'source': 'git-subdir', 'url': 'https://github.com/org/repo.git', 'path': 'plugins/plugin-name'}\n" + "- Zip archive on any https host (e.g. S3): " + "{'source': 'archive', 'url': 'https://bucket.s3.amazonaws.com/plugin.zip', 'sha256': ''}" ), ) version: str | None = Field("1.0.0", description="Semantic version") @@ -46,7 +48,7 @@ class RegisterPluginRequest(PluginSpec): Request body for registering a plugin in the marketplace. LiteLLM acts as a registry/discovery layer. Plugins are hosted on - GitHub/GitLab/Bitbucket and referenced by their git source. + GitHub/GitLab/Bitbucket or as a zip archive on any https host and referenced by their source. """ name: str = Field( @@ -76,7 +78,7 @@ class PluginResponse(BaseModel): name: str = Field(..., description="Plugin name") version: str | None = Field(None, description="Plugin version") description: str | None = Field(None, description="Plugin description") - source: dict[str, str] = Field(..., description="Git source reference") + source: dict[str, str] = Field(..., description="Plugin source reference") enabled: bool = Field(..., description="Whether plugin is enabled") diff --git a/tests/test_litellm/proxy/anthropic_endpoints/test_claude_code_marketplace.py b/tests/test_litellm/proxy/anthropic_endpoints/test_claude_code_marketplace.py index 99b09f48fe5..aa2da262bae 100644 --- a/tests/test_litellm/proxy/anthropic_endpoints/test_claude_code_marketplace.py +++ b/tests/test_litellm/proxy/anthropic_endpoints/test_claude_code_marketplace.py @@ -1,7 +1,7 @@ """ Unit tests for claude_code_marketplace.py source validation. -Covers the git-subdir source type added alongside the existing github and url types. +Covers the git-subdir and archive source types added alongside the existing github and url types. """ import json @@ -87,6 +87,12 @@ _GIT_SUBDIR_SOURCE = { "path": "plugins/my-plugin", } +_ARCHIVE_SOURCE = { + "source": "archive", + "url": "https://skills-bucket.s3.us-east-1.amazonaws.com/plugins/s3-skill-1.0.0.zip", + "sha256": "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855", +} + @pytest.fixture(autouse=True) def _patch_proxy_globals(monkeypatch): @@ -377,6 +383,75 @@ async def test_register_plugin_unknown_source_type(): assert exc_info.value.status_code == 400 assert "git-subdir" in exc_info.value.detail["error"] + assert "archive" in exc_info.value.detail["error"] + + +@pytest.mark.asyncio +async def test_archive_source_registers_and_is_served_verbatim_in_marketplace(): + response = await register_plugin( + request=RegisterPluginRequest(name="s3-skill", source=_ARCHIVE_SOURCE), + user_api_key_dict=_USER, + ) + + assert response.action == "created" + assert response.plugin.source == _ARCHIVE_SOURCE + + marketplace = json.loads((await get_marketplace()).body) + assert marketplace["plugins"] == [{"name": "s3-skill", "source": _ARCHIVE_SOURCE, "version": "1.0.0"}] + + +@pytest.mark.asyncio +async def test_archive_source_without_sha256_is_accepted(): + source = {"source": "archive", "url": "https://artifacts.example.com/plugin.zip"} + + response = await register_plugin( + request=RegisterPluginRequest(name="unpinned-skill", source=source), + user_api_key_dict=_USER, + ) + + assert response.plugin.source == source + + +@pytest.mark.asyncio +async def test_update_plugin_to_archive_source(): + name = "my-monorepo-plugin" + await register_plugin( + request=RegisterPluginRequest(name=name, source=_GIT_SUBDIR_SOURCE), + user_api_key_dict=_USER, + ) + + response = await update_plugin( + plugin_name=name, + request=UpdatePluginRequest(source=_ARCHIVE_SOURCE), + user_api_key_dict=_USER, + ) + + assert response.action == "updated" + assert (await _read_stored_manifest(name))["source"] == _ARCHIVE_SOURCE + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "source, expected_fragment", + [ + ({"source": "archive"}, "url"), + ({"source": "archive", "url": ""}, "url"), + ({"source": "archive", "url": "http://artifacts.example.com/plugin.zip"}, "https"), + ({"source": "archive", "url": "s3://skills-bucket/plugin.zip"}, "https"), + ({"source": "archive", "url": "https://"}, "https"), + ({"source": "archive", "url": "https:///plugin.zip"}, "https"), + ({"source": "archive", "url": "https://[::1/plugin.zip"}, "https"), + ({"source": "archive", "url": "https://artifacts.example.com/plugin.zip", "sha256": "a" * 63}, "sha256"), + ({"source": "archive", "url": "https://artifacts.example.com/plugin.zip", "sha256": "a" * 65}, "sha256"), + ({"source": "archive", "url": "https://artifacts.example.com/plugin.zip", "sha256": "g" * 64}, "sha256"), + ], +) +async def test_register_plugin_archive_rejects_malformed_source(source, expected_fragment): + with pytest.raises(HTTPException) as exc_info: + await register_plugin(request=RegisterPluginRequest(name="bad-plugin", source=source), user_api_key_dict=_USER) + + assert exc_info.value.status_code == 400 + assert expected_fragment in exc_info.value.detail["error"] @pytest.mark.asyncio diff --git a/ui/litellm-dashboard/src/app/(dashboard)/skills/_components/add_plugin_form.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/skills/_components/add_plugin_form.test.tsx index 1058d4d0466..6507855d0c7 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/skills/_components/add_plugin_form.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/skills/_components/add_plugin_form.test.tsx @@ -20,21 +20,25 @@ const DEFAULT_PROPS = { onSuccess: vi.fn(), }; -const URL_PLACEHOLDER = "https://github.com/org/repo or https://gitlab.com/org/repo"; +const URL_PLACEHOLDER = "https://github.com/org/repo or https://bucket.s3.amazonaws.com/my-skill.zip"; const SUBPATH_PLACEHOLDER = "plugins/my-skill"; +const SHA256_PLACEHOLDER = "64 hex characters"; +const S3_ZIP_URL = "https://skills-bucket.s3.us-east-1.amazonaws.com/plugins/s3-skill-1.0.0.zip"; +const DIGEST = "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855"; describe("AddPluginForm", () => { beforeEach(() => { vi.clearAllMocks(); }); - it("renders the host-agnostic repository URL input and subfolder field", () => { + it("renders the host-agnostic source URL input and subfolder field, hiding the digest until a zip is entered", () => { renderWithProviders(); - expect(screen.getByText("Repository URL")).toBeInTheDocument(); + expect(screen.getByText("Source URL")).toBeInTheDocument(); expect(screen.getByPlaceholderText(URL_PLACEHOLDER)).toBeInTheDocument(); expect(screen.getByText("Subfolder path (Optional)")).toBeInTheDocument(); expect(screen.getByPlaceholderText(SUBPATH_PLACEHOLDER)).toBeInTheDocument(); + expect(screen.queryByPlaceholderText(SHA256_PLACEHOLDER)).not.toBeInTheDocument(); }); it("shows GitHub repo preview for a plain repo URL", async () => { @@ -255,6 +259,85 @@ describe("AddPluginForm", () => { }); }); + it("shows a zip archive preview, disables the subfolder field, and reveals the digest field", async () => { + renderWithProviders(); + + await typeUrl(S3_ZIP_URL); + + expect(await screen.findByText(/Zip archive/)).toBeInTheDocument(); + expect(screen.getByPlaceholderText(SUBPATH_PLACEHOLDER)).toBeDisabled(); + expect(screen.getByText("A zip archive is installed as a whole, so this field is disabled")).toBeInTheDocument(); + expect(screen.getByPlaceholderText(SHA256_PLACEHOLDER)).toBeInTheDocument(); + expect((screen.getByPlaceholderText("my-skill") as HTMLInputElement).value).toBe("s3-skill-1-0-0"); + }); + + it("submits an archive source without a digest when the field is left empty", async () => { + renderWithProviders(); + + await typeUrl(S3_ZIP_URL); + await submit(); + + await waitFor(() => { + expect(mockRegister).toHaveBeenCalledWith( + "sk-test", + expect.objectContaining({ source: { source: "archive", url: S3_ZIP_URL } }), + ); + }); + }); + + it("submits an archive source pinned to the lowercased digest", async () => { + renderWithProviders(); + + await typeUrl(S3_ZIP_URL); + await act(async () => { + fireEvent.change(screen.getByPlaceholderText(SHA256_PLACEHOLDER), { + target: { value: ` ${DIGEST.toUpperCase()} ` }, + }); + }); + await submit(); + + await waitFor(() => { + expect(mockRegister).toHaveBeenCalledWith( + "sk-test", + expect.objectContaining({ source: { source: "archive", url: S3_ZIP_URL, sha256: DIGEST } }), + ); + }); + }); + + it("drops the digest once the archive URL changes so a stale checksum is never sent for a new file", async () => { + renderWithProviders(); + + await typeUrl(S3_ZIP_URL); + await act(async () => { + fireEvent.change(screen.getByPlaceholderText(SHA256_PLACEHOLDER), { target: { value: DIGEST } }); + }); + const otherZipUrl = S3_ZIP_URL.replace("1.0.0", "1.1.0"); + await typeUrl(otherZipUrl); + + expect(screen.getByPlaceholderText(SHA256_PLACEHOLDER)).toHaveValue(""); + await submit(); + + await waitFor(() => { + expect(mockRegister).toHaveBeenCalledWith( + "sk-test", + expect.objectContaining({ source: { source: "archive", url: otherZipUrl } }), + ); + }); + }); + + it("blocks submission and shows the digest error for a malformed sha256", async () => { + renderWithProviders(); + + await typeUrl(S3_ZIP_URL); + await act(async () => { + fireEvent.change(screen.getByPlaceholderText(SHA256_PLACEHOLDER), { target: { value: "not-a-digest" } }); + }); + await submit(); + + expect(await screen.findByText("SHA-256 must be a 64-character hex digest")).toBeInTheDocument(); + expect(mockRegister).not.toHaveBeenCalled(); + }); + it("surfaces the backend error message when registration fails", async () => { mockRegister.mockRejectedValueOnce(new Error("Plugin 'claude-code' already exists")); renderWithProviders(); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/skills/_components/add_plugin_form.tsx b/ui/litellm-dashboard/src/app/(dashboard)/skills/_components/add_plugin_form.tsx index 70669ffc0cc..e420038b124 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/skills/_components/add_plugin_form.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/skills/_components/add_plugin_form.tsx @@ -26,6 +26,7 @@ import { parseKeywords, parseSkillSource, isValidSubPath, + isValidSha256, SkillSourcePreview, } from "@/components/claude_code_plugins/helpers"; import { PluginAuthor, PluginSource, SkillRegisterRequest } from "@/components/claude_code_plugins/types"; @@ -39,13 +40,14 @@ interface AddPluginFormProps { } const addPluginShape = { - skillUrl: z.string().min(1, "Please enter a repository URL"), + skillUrl: z.string().min(1, "Please enter a repository or zip archive URL"), subPath: z .string() .refine( (value) => !value || isValidSubPath(value), "Subfolder must be a relative path like plugins/my-skill (letters, numbers, dots, hyphens, underscores)", ), + sha256: z.string().refine(isValidSha256, "SHA-256 must be a 64-character hex digest"), name: z .string() .min(1, "Please enter skill name") @@ -69,6 +71,7 @@ type AddPluginFormValues = z.infer; const EMPTY_VALUES: AddPluginFormValues = { skillUrl: "", subPath: "", + sha256: "", name: "", domain: "", namespace: "", @@ -89,11 +92,19 @@ const buildAuthor = (values: AddPluginFormValues): PluginAuthor | undefined => { return email ? { name, email } : { name }; }; +const archiveUrlOf = (preview: SkillSourcePreview | null): string | undefined => + preview?.parsed.source === "archive" ? preview.parsed.url : undefined; + +const withArchiveDigest = (source: PluginSource, sha256: string): PluginSource => { + const digest = sha256.trim(); + return source.source === "archive" && digest ? { ...source, sha256: digest.toLowerCase() } : source; +}; + const buildRegisterRequest = (values: AddPluginFormValues, source: PluginSource): SkillRegisterRequest => { const author = buildAuthor(values); return { name: values.name.trim(), - source, + source: withArchiveDigest(source, values.sha256), ...(values.version ? { version: values.version.trim() } : {}), ...(values.description ? { description: values.description.trim() } : {}), ...(author ? { author } : {}), @@ -115,6 +126,16 @@ const PREDEFINED_CATEGORIES = [ "Documentation", ]; +const SUB_PATH_LOCK_REASON = { + "git-subdir": "The URL already points to a subfolder, so this field is disabled", + archive: "A zip archive is installed as a whole, so this field is disabled", +} as const; + +type SubPathLock = keyof typeof SUB_PATH_LOCK_REASON; + +const subPathLockFor = (source: PluginSource["source"] | undefined): SubPathLock | null => + source === "git-subdir" || source === "archive" ? source : null; + const labelWithHint = (label: string, hint: string): React.ReactNode => ( <> {label} @@ -129,15 +150,18 @@ const AddPluginForm: React.FC = ({ visible, onClose, accessT const form = useZodForm(addPluginSchema, { defaultValues: EMPTY_VALUES }); const [isSubmitting, setIsSubmitting] = useState(false); const [urlPreview, setUrlPreview] = useState(null); - const [urlEncodesSubdir, setUrlEncodesSubdir] = useState(false); + const [subPathLock, setSubPathLock] = useState(null); const recomputePreview = (skillUrl: string, subPath: string) => { - const encodesSubdir = parseSkillSource(skillUrl)?.parsed.source === "git-subdir"; - setUrlEncodesSubdir(encodesSubdir); - if (encodesSubdir && form.getValues("subPath")) { + const lock = subPathLockFor(parseSkillSource(skillUrl)?.parsed.source); + setSubPathLock(lock); + if (lock && form.getValues("subPath")) { form.setValue("subPath", ""); } - const preview = parseSkillSource(skillUrl, encodesSubdir ? undefined : subPath); + const preview = parseSkillSource(skillUrl, lock ? undefined : subPath); + if (archiveUrlOf(preview) !== archiveUrlOf(urlPreview) && form.getValues("sha256")) { + form.setValue("sha256", ""); + } setUrlPreview(preview); if (preview && !form.getValues("name")) { form.setValue("name", preview.suggestedName); @@ -151,7 +175,7 @@ const AddPluginForm: React.FC = ({ visible, onClose, accessT } if (!urlPreview) { - toast.error("Please enter a valid repository URL"); + toast.error("Please enter a valid repository or zip archive URL"); return; } @@ -176,7 +200,7 @@ const AddPluginForm: React.FC = ({ visible, onClose, accessT toast.success("Skill registered successfully"); form.reset(EMPTY_VALUES); setUrlPreview(null); - setUrlEncodesSubdir(false); + setSubPathLock(null); onSuccess(); onClose(); } catch (error) { @@ -190,7 +214,7 @@ const AddPluginForm: React.FC = ({ visible, onClose, accessT const handleCancel = () => { form.reset(EMPTY_VALUES); setUrlPreview(null); - setUrlEncodesSubdir(false); + setSubPathLock(null); onClose(); }; @@ -207,15 +231,15 @@ const AddPluginForm: React.FC = ({ visible, onClose, accessT control={form.control} name="skillUrl" label={labelWithHint( - "Repository URL", - "Paste an HTTPS git repository URL from GitHub, GitLab, Bitbucket, or a self-hosted host. E.g. github.com/org/repo, gitlab.com/org/repo, or github.com/org/repo/tree/main/my-skill", + "Source URL", + "Paste an HTTPS git repository URL from GitHub, GitLab, Bitbucket, or a self-hosted host (e.g. github.com/org/repo or github.com/org/repo/tree/main/my-skill), or an HTTPS link to a .zip archive of the skill hosted on S3 or any static file server.", )} > {({ ref, onChange, ...field }) => ( { onChange(event); @@ -232,9 +256,7 @@ const AddPluginForm: React.FC = ({ visible, onClose, accessT "Subfolder path (Optional)", "Path within the repository where the skill lives (e.g., plugins/my-skill). Leave empty if the skill is at the repo root.", )} - description={ - urlEncodesSubdir ? "The URL already points to a subfolder, so this field is disabled" : undefined - } + description={subPathLock ? SUB_PATH_LOCK_REASON[subPathLock] : undefined} > {({ ref, onChange, ...field }) => ( = ({ visible, onClose, accessT onChange(event); recomputePreview(form.getValues("skillUrl"), event.target.value); }} - disabled={urlEncodesSubdir} + disabled={subPathLock !== null} /> )} + {urlPreview?.parsed.source === "archive" && ( + + {({ ref, ...field }) => ( + + )} + + )} + {urlPreview && (
Detected: {urlPreview.label} diff --git a/ui/litellm-dashboard/src/components/AIHub/SkillHubTableColumns.tsx b/ui/litellm-dashboard/src/components/AIHub/SkillHubTableColumns.tsx index 2a1530cc352..1104e2c7485 100644 --- a/ui/litellm-dashboard/src/components/AIHub/SkillHubTableColumns.tsx +++ b/ui/litellm-dashboard/src/components/AIHub/SkillHubTableColumns.tsx @@ -26,7 +26,7 @@ function getSkillSourceLink(skill: Plugin): { url: string; label: string } | nul const url = src.path ? `${src.url}/tree/main/${src.path}` : src.url; return { url, label: url.replace("https://github.com/", "") }; } - if (src?.source === "url" && src.url) { + if ((src?.source === "url" || src?.source === "archive") && src.url) { return { url: src.url, label: src.url.replace(/^https?:\/\//, "") }; } return null; diff --git a/ui/litellm-dashboard/src/components/claude_code_plugins/helpers.test.ts b/ui/litellm-dashboard/src/components/claude_code_plugins/helpers.test.ts index 1d07ca09df5..4713fcbc55b 100644 --- a/ui/litellm-dashboard/src/components/claude_code_plugins/helpers.test.ts +++ b/ui/litellm-dashboard/src/components/claude_code_plugins/helpers.test.ts @@ -17,6 +17,7 @@ import { formatKeywords, parseSkillSource, isValidSubPath, + isValidSha256, buildMarketplaceSettingsSnippet, } from "./helpers"; import { MarketplacePluginEntry } from "./types"; @@ -114,6 +115,12 @@ describe("getSourceDisplayText", () => { ); }); + it("shows the archive url for an archive source", () => { + expect(getSourceDisplayText({ source: "archive", url: "https://bucket.s3.amazonaws.com/skill.zip" })).toBe( + "https://bucket.s3.amazonaws.com/skill.zip", + ); + }); + it("returns unknown for missing data", () => { expect(getSourceDisplayText({ source: "github" })).toBe("Unknown source"); }); @@ -140,6 +147,12 @@ describe("getSourceLink", () => { ); }); + it("returns the archive url for an archive source", () => { + expect(getSourceLink({ source: "archive", url: "https://bucket.s3.amazonaws.com/skill.zip" })).toBe( + "https://bucket.s3.amazonaws.com/skill.zip", + ); + }); + it("returns null when no repo or url", () => { expect(getSourceLink({ source: "github" })).toBeNull(); }); @@ -329,6 +342,20 @@ describe("isValidUrl", () => { }); }); +describe("isValidSha256", () => { + it("accepts an empty digest and a 64-character hex digest in either case", () => { + expect(isValidSha256("")).toBe(true); + expect(isValidSha256("a".repeat(64))).toBe(true); + expect(isValidSha256(" " + "ABCDEF0123456789".repeat(4) + " ")).toBe(true); + }); + + it("rejects wrong length and non-hex digests", () => { + expect(isValidSha256("a".repeat(63))).toBe(false); + expect(isValidSha256("a".repeat(65))).toBe(false); + expect(isValidSha256("g".repeat(64))).toBe(false); + }); +}); + describe("parseKeywords", () => { it("splits comma-separated keywords", () => { expect(parseKeywords("a, b, c")).toEqual(["a", "b", "c"]); @@ -445,6 +472,39 @@ describe("parseSkillSource", () => { expect(parseSkillSource("not a url")).toBeNull(); }); + it("parses an S3 zip URL into an archive source and names the skill after the file", () => { + expect(parseSkillSource("https://skills-bucket.s3.us-east-1.amazonaws.com/plugins/My_Skill-1.0.0.zip")).toEqual({ + parsed: { source: "archive", url: "https://skills-bucket.s3.us-east-1.amazonaws.com/plugins/My_Skill-1.0.0.zip" }, + label: "Zip archive — skills-bucket.s3.us-east-1.amazonaws.com/plugins/My_Skill-1.0.0.zip", + suggestedName: "my-skill-1-0-0", + }); + }); + + it("keeps the query string of a zip URL so versioned or signed object links still resolve", () => { + expect(parseSkillSource("https://bucket.s3.amazonaws.com/skill.ZIP?versionId=abc")?.parsed).toEqual({ + source: "archive", + url: "https://bucket.s3.amazonaws.com/skill.ZIP?versionId=abc", + }); + }); + + it("ignores the subfolder for a zip URL since the archive is installed whole", () => { + expect(parseSkillSource("https://artifacts.example.com/skill.zip", "plugins/x")?.parsed).toEqual({ + source: "archive", + url: "https://artifacts.example.com/skill.zip", + }); + }); + + it("rejects a plain http zip URL", () => { + expect(parseSkillSource("http://artifacts.example.com/skill.zip")).toBeNull(); + }); + + it("treats a github zip download URL as an archive rather than a repo path", () => { + expect(parseSkillSource("https://github.com/org/repo/releases/download/v1/skill.zip")?.parsed).toEqual({ + source: "archive", + url: "https://github.com/org/repo/releases/download/v1/skill.zip", + }); + }); + it("suggests a kebab-friendly name from the last path segment", () => { expect(parseSkillSource("github.com/org/my-awesome-skill")?.suggestedName).toBe("my-awesome-skill"); expect(parseSkillSource("github.com/org/repo/tree/main/plugins/cool-skill")?.suggestedName).toBe("cool-skill"); diff --git a/ui/litellm-dashboard/src/components/claude_code_plugins/helpers.ts b/ui/litellm-dashboard/src/components/claude_code_plugins/helpers.ts index 5b3d4f20c1e..6f673d67f87 100644 --- a/ui/litellm-dashboard/src/components/claude_code_plugins/helpers.ts +++ b/ui/litellm-dashboard/src/components/claude_code_plugins/helpers.ts @@ -23,6 +23,12 @@ const GITHUB_HOST = "github.com"; const SKILL_FILE_EXTENSION_REGEX = /\.(md|markdown|txt|json|ya?ml|toml)$/i; +const ZIP_ARCHIVE_REGEX = /\.zip$/i; + +export const SHA256_REGEX = /^[0-9a-fA-F]{64}$/; + +export const isValidSha256 = (digest: string): boolean => digest.trim() === "" || SHA256_REGEX.test(digest.trim()); + // WHATWG normalizes obfuscated IPv4 (e.g. 2130706433, 0x7f.0.0.1) to dotted-decimal, so this // catches every IPv4 form; bracketed IPv6 is rejected separately. const IPV4_HOST_REGEX = /^\d{1,3}(\.\d{1,3}){3}$/; @@ -160,16 +166,26 @@ const parseRawGitSource = (url: URL, subPath?: string): SkillSourcePreview | nul }; }; +const parseArchiveSource = (url: URL): SkillSourcePreview => ({ + parsed: { source: "archive", url: url.href }, + label: `Zip archive — ${url.host}${url.pathname}`, + suggestedName: toKebabCase(lastSegment(url.pathname).replace(ZIP_ARCHIVE_REGEX, "")), +}); + /** - * Parse any git-accessible repository URL into a registerable skill source. - * GitHub URLs keep their `github`/`git-subdir` shorthand; every other host is - * treated as a raw repo URL, with an optional subfolder turning it into git-subdir. + * Parse any git-accessible repository URL or https zip archive URL into a registerable skill + * source. A `.zip` path is an `archive` source (S3, Artifactory, any static host). GitHub URLs + * keep their `github`/`git-subdir` shorthand; every other host is treated as a raw repo URL, + * with an optional subfolder turning it into git-subdir. */ export const parseSkillSource = (rawUrl: string, subPath?: string): SkillSourcePreview | null => { const url = parseRepoUrl(rawUrl); if (!url) { return null; } + if (ZIP_ARCHIVE_REGEX.test(url.pathname)) { + return parseArchiveSource(url); + } if (url.hostname.replace(/^www\./, "") === GITHUB_HOST) { return parseGitHubSource(url, subPath); } @@ -245,7 +261,7 @@ export const getSourceDisplayText = (source: PluginSource): string => { if (source.source === "git-subdir" && source.url && source.path) { return `${source.url} @ ${source.path}`; } - if (source.source === "url" && source.url) { + if ((source.source === "url" || source.source === "archive") && source.url) { return source.url; } return "Unknown source"; @@ -258,10 +274,8 @@ export const getSourceLink = (source: PluginSource): string | null => { if (source.source === "github" && source.repo) { return `https://github.com/${source.repo}`; } - if ((source.source === "url" || source.source === "git-subdir") && source.url) { - return source.url; - } - return null; + const linksToUrl = source.source === "url" || source.source === "git-subdir" || source.source === "archive"; + return linksToUrl && source.url ? source.url : null; }; /** diff --git a/ui/litellm-dashboard/src/components/claude_code_plugins/skill_detail.tsx b/ui/litellm-dashboard/src/components/claude_code_plugins/skill_detail.tsx index 7cbbb08b623..b37204c6073 100644 --- a/ui/litellm-dashboard/src/components/claude_code_plugins/skill_detail.tsx +++ b/ui/litellm-dashboard/src/components/claude_code_plugins/skill_detail.tsx @@ -26,7 +26,7 @@ const SkillDetail: React.FC = ({ skill, onBack }) => { const src = skill.source; if (src.source === "github" && src.repo) return `https://github.com/${src.repo}`; if (src.source === "git-subdir" && src.url) return src.path ? `${src.url}/tree/main/${src.path}` : src.url; - if (src.source === "url" && src.url) return src.url; + if ((src.source === "url" || src.source === "archive") && src.url) return src.url; return null; })(); diff --git a/ui/litellm-dashboard/src/components/claude_code_plugins/types.ts b/ui/litellm-dashboard/src/components/claude_code_plugins/types.ts index d16c880749b..98f80fd2959 100644 --- a/ui/litellm-dashboard/src/components/claude_code_plugins/types.ts +++ b/ui/litellm-dashboard/src/components/claude_code_plugins/types.ts @@ -8,10 +8,11 @@ import type { components } from "@/lib/http/schema"; // Kept hand-written: the backend types `source` as Dict[str, str], so the generated type is a // loose string map; this discriminant union is what the parser and display helpers rely on. export interface PluginSource { - source: "github" | "url" | "git-subdir"; + source: "github" | "url" | "git-subdir" | "archive"; repo?: string; // Format: "org/repo" for GitHub url?: string; // Full URL for other sources path?: string; // Subdirectory path for git-subdir + sha256?: string; } export type PluginAuthor = components["schemas"]["PluginAuthor"]; diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index ab2c1f27a59..c7a760e1eee 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -2152,8 +2152,8 @@ export interface paths { * @description Register a new plugin in the LiteLLM marketplace. * * LiteLLM acts as a registry/discovery layer. Plugins are hosted on - * GitHub/GitLab/Bitbucket. Claude Code will clone from the git source - * when users install. + * GitHub/GitLab/Bitbucket or as a zip archive on any https host (e.g. S3). + * Claude Code clones the git source or downloads the archive when users install. * * This endpoint is create-only and never overwrites. If a plugin with * the same name already exists it returns 409 Conflict; use @@ -2163,7 +2163,7 @@ export interface paths { * * Parameters: * - name: Plugin name (kebab-case) - * - source: Git source reference (github, url, or git-subdir format) + * - source: Plugin source reference (github, url, git-subdir, or archive format) * - version: Semantic version (optional) * - description: Plugin description (optional) * - author: Author information (optional) @@ -2229,7 +2229,7 @@ export interface paths { * * Parameters: * - plugin_name: Name of the plugin to update (path parameter) - * - source: Git source reference (github, url, or git-subdir format) + * - source: Plugin source reference (github, url, git-subdir, or archive format) * - version: Semantic version (optional) * - description: Plugin description (optional) * - author: Author information (optional) @@ -33638,7 +33638,7 @@ export interface components { name: string; /** * Source - * @description Git source reference + * @description Plugin source reference */ source: { [key: string]: string; @@ -34887,7 +34887,7 @@ export interface components { * @description Request body for registering a plugin in the marketplace. * * LiteLLM acts as a registry/discovery layer. Plugins are hosted on - * GitHub/GitLab/Bitbucket and referenced by their git source. + * GitHub/GitLab/Bitbucket or as a zip archive on any https host and referenced by their source. */ RegisterPluginRequest: { /** @description Plugin author */ @@ -34929,10 +34929,11 @@ export interface components { namespace?: string | null; /** * Source - * @description Git source reference. Supported formats: + * @description Plugin source reference. Supported formats: * - GitHub: {'source': 'github', 'repo': 'org/repo'} * - Git URL: {'source': 'url', 'url': 'https://github.com/org/repo.git'} * - Git Subdir: {'source': 'git-subdir', 'url': 'https://github.com/org/repo.git', 'path': 'plugins/plugin-name'} + * - Zip archive on any https host (e.g. S3): {'source': 'archive', 'url': 'https://bucket.s3.amazonaws.com/plugin.zip', 'sha256': ''} */ source: { [key: string]: string; @@ -38211,10 +38212,11 @@ export interface components { namespace?: string | null; /** * Source - * @description Git source reference. Supported formats: + * @description Plugin source reference. Supported formats: * - GitHub: {'source': 'github', 'repo': 'org/repo'} * - Git URL: {'source': 'url', 'url': 'https://github.com/org/repo.git'} * - Git Subdir: {'source': 'git-subdir', 'url': 'https://github.com/org/repo.git', 'path': 'plugins/plugin-name'} + * - Zip archive on any https host (e.g. S3): {'source': 'archive', 'url': 'https://bucket.s3.amazonaws.com/plugin.zip', 'sha256': ''} */ source: { [key: string]: string; From 6cc13e07a68f1e822c1acbb7bec808c45baec8e2 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 10 Sep 2026 14:26:09 -0700 Subject: [PATCH 21/21] feat(proxy): granular key/team access control for Claude Code marketplace plugins (#40518) Adds object_permission.skills to keys and teams, enforces it on /claude-code/marketplace.json?key=, /claude-code/plugins and /claude-code/plugins/{name}, and exposes an Allowed Skills selector in the key and team create/edit forms of the Admin UI Co-authored-by: yassin Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../migration.sql | 1 + .../litellm_proxy_extras/schema.prisma | 1 + litellm/models/object_permission.py | 1 + litellm/proxy/_lazy_openapi_snapshot.json | 30 ++++- litellm/proxy/_types.py | 4 + .../claude_code_marketplace.py | 43 +++++-- .../claude_code_skill_access.py | 67 +++++++++++ litellm/proxy/schema.prisma | 1 + .../object_permission_repository.py | 6 + litellm/types/object_permission.py | 3 +- schema.prisma | 1 + .../test_claude_code_marketplace.py | 2 +- .../test_claude_code_marketplace.py | 97 +++++++++++++++- .../test_claude_code_skill_access.py | 79 +++++++++++++ .../proxy/auth/test_route_checks.py | 14 +++ .../test_customer_endpoints.py | 1 + .../test_object_permission_utils.py | 23 ++++ .../src/components/Teams.test.tsx | 29 +++++ ui/litellm-dashboard/src/components/Teams.tsx | 46 ++++++++ .../components/object_permissions_view.tsx | 11 ++ .../organisms/createKeyPayload.test.ts | 12 ++ .../components/organisms/createKeyPayload.ts | 5 + .../create_key_button.integration.test.tsx | 24 ++++ .../organisms/create_key_button.tsx | 31 +++++ .../components/skills/SkillSelector.test.tsx | 54 +++++++++ .../src/components/skills/SkillSelector.tsx | 109 ++++++++++++++++++ .../src/components/team/TeamInfo.test.tsx | 43 +++++++ .../src/components/team/TeamInfo.tsx | 26 +++++ .../templates/KeyEditViewControls.tsx | 34 +++++- .../KeyInfoView.handleKeyUpdate.test.tsx | 39 +++++++ .../components/templates/keyEditFormValues.ts | 4 + .../templates/key_edit_view.test.tsx | 40 +++++++ .../components/templates/key_edit_view.tsx | 15 +-- .../components/templates/key_info_view.tsx | 8 ++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 22 +++- 35 files changed, 892 insertions(+), 34 deletions(-) create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20260910000000_skills_on_object_permission/migration.sql create mode 100644 litellm/proxy/anthropic_endpoints/claude_code_endpoints/claude_code_skill_access.py create mode 100644 tests/test_litellm/proxy/anthropic_endpoints/test_claude_code_skill_access.py create mode 100644 ui/litellm-dashboard/src/components/skills/SkillSelector.test.tsx create mode 100644 ui/litellm-dashboard/src/components/skills/SkillSelector.tsx diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260910000000_skills_on_object_permission/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260910000000_skills_on_object_permission/migration.sql new file mode 100644 index 00000000000..c982bc38a69 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260910000000_skills_on_object_permission/migration.sql @@ -0,0 +1 @@ +ALTER TABLE "LiteLLM_ObjectPermissionTable" ADD COLUMN IF NOT EXISTS "skills" TEXT[] DEFAULT ARRAY[]::TEXT[]; diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 05c5aad9303..817df082d8c 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -282,6 +282,7 @@ model LiteLLM_ObjectPermissionTable { mcp_toolsets String[] @default([]) // Toolset IDs granted to this key/team/user search_tools String[] @default([]) // search_tool_name values this key/team/user may call mcp_tool_search_enabled Boolean? + skills String[] @default([]) // Claude Code plugin names granted to this key/team beyond the public (enabled) set teams LiteLLM_TeamTable[] projects LiteLLM_ProjectTable[] verification_tokens LiteLLM_VerificationToken[] diff --git a/litellm/models/object_permission.py b/litellm/models/object_permission.py index a09d50ddc33..f178c5ad47f 100644 --- a/litellm/models/object_permission.py +++ b/litellm/models/object_permission.py @@ -23,3 +23,4 @@ class LiteLLM_ObjectPermissionTable(LiteLLMPydanticObjectBase): blocked_tools: list[str] | None = [] search_tools: list[str] | None = [] mcp_tool_search_enabled: bool | None = None + skills: list[str] | None = None diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 3f060f1f5f6..f0af17ab818 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -5721,8 +5721,26 @@ "paths": { "/claude-code/marketplace.json": { "get": { - "description": "Serve marketplace.json for Claude Code plugin discovery.\n\nThis endpoint is accessed by Claude Code CLI when users run:\n- claude plugin marketplace add \n- claude plugin install @\n\nReturns:\n Marketplace catalog with list of available plugins and their git sources.\n\nExample:\n ```bash\n claude plugin marketplace add http://localhost:4000/claude-code/marketplace.json\n claude plugin install my-plugin@litellm\n ```", + "description": "Serve marketplace.json for Claude Code plugin discovery.\n\nThis endpoint is accessed by Claude Code CLI when users run:\n- claude plugin marketplace add \n- claude plugin install @\n\nWithout `key` the catalog holds the enabled (public) plugins. With `?key=sk-...`\nthe key is authenticated and the catalog also holds the disabled plugins granted\nto it through `object_permission.skills` on the key or its team.\n\nReturns:\n Marketplace catalog with list of available plugins and their git sources.\n\nExample:\n ```bash\n claude plugin marketplace add http://localhost:4000/claude-code/marketplace.json\n claude plugin marketplace add \"http://localhost:4000/claude-code/marketplace.json?key=sk-...\"\n claude plugin install my-plugin@litellm\n ```", "operationId": "get_marketplace_claude_code_marketplace_json_get", + "parameters": [ + { + "in": "query", + "name": "key", + "required": false, + "schema": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Key" + } + } + ], "responses": { "200": { "content": { @@ -5731,6 +5749,16 @@ } }, "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" } }, "summary": "Get Marketplace", diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index f496140e905..23dc8237160 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -497,6 +497,9 @@ class LiteLLMRoutes(enum.Enum): "/v1/messages/count_tokens", "/v1/skills", "/v1/skills/{skill_id}", + "/claude-code/marketplace.json", + "/claude-code/plugins", + "/claude-code/plugins/{plugin_name}", ] # MCP tool-call / passthrough routes — data-plane. Gated by DISABLE_LLM_API_ENDPOINTS. @@ -1124,6 +1127,7 @@ class LiteLLM_ObjectPermissionBase(LiteLLMPydanticObjectBase): models: list[str] | None = None search_tools: list[str] | None = None mcp_tool_search_enabled: bool | None = None + skills: list[str] | None = None from litellm.models.team import BudgetLimitEntry as BudgetLimitEntry # noqa: E402 diff --git a/litellm/proxy/anthropic_endpoints/claude_code_endpoints/claude_code_marketplace.py b/litellm/proxy/anthropic_endpoints/claude_code_endpoints/claude_code_marketplace.py index cdeb1c6a111..7c6a4571948 100644 --- a/litellm/proxy/anthropic_endpoints/claude_code_endpoints/claude_code_marketplace.py +++ b/litellm/proxy/anthropic_endpoints/claude_code_endpoints/claude_code_marketplace.py @@ -7,10 +7,10 @@ Actual plugin files are hosted on GitHub/GitLab/Bitbucket or as a zip archive on any HTTPS host (S3, Artifactory, a static file server). Endpoints: -/claude-code/marketplace.json - GET - List plugins for Claude Code discovery (unauthenticated) +/claude-code/marketplace.json - GET - List plugins for Claude Code discovery (unauthenticated; `?key=` adds the key's granted skills) /claude-code/plugins - POST - Register a new plugin (create-only, proxy admin only) -/claude-code/plugins - GET - List plugins (any authenticated key) -/claude-code/plugins/{name} - GET - Get plugin details (any authenticated key) +/claude-code/plugins - GET - List plugins visible to the key (enabled, plus granted disabled ones) +/claude-code/plugins/{name} - GET - Get plugin details (403 on a disabled plugin the key is not granted) /claude-code/plugins/{name} - PUT - Update an existing plugin (proxy admin only) /claude-code/plugins/{name}/enable - POST - Enable a plugin (proxy admin only) /claude-code/plugins/{name}/disable - POST - Disable a plugin (proxy admin only) @@ -24,11 +24,15 @@ from datetime import datetime, timezone from typing import Annotated, Final, Protocol, TypedDict from urllib.parse import urlsplit -from fastapi import APIRouter, Depends, HTTPException +from fastapi import APIRouter, Depends, HTTPException, Request from fastapi.responses import JSONResponse from litellm._logging import verbose_proxy_logger -from litellm.proxy._types import CommonProxyErrors, UserAPIKeyAuth +from litellm.proxy._types import CommonProxyErrors, ProxyException, UserAPIKeyAuth +from litellm.proxy.anthropic_endpoints.claude_code_endpoints.claude_code_skill_access import ( + SkillVisibility, + skill_visibility, +) from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.resource_ownership import is_proxy_admin from litellm.repositories.table_repositories import ClaudeCodePluginRepository @@ -84,7 +88,7 @@ async def _get_prisma_client() -> object: "/claude-code/marketplace.json", tags=["Claude Code Marketplace"], ) -async def get_marketplace(): +async def get_marketplace(request: Request, key: str | None = None): """ Serve marketplace.json for Claude Code plugin discovery. @@ -92,24 +96,35 @@ async def get_marketplace(): - claude plugin marketplace add - claude plugin install @ + Without `key` the catalog holds the enabled (public) plugins. With `?key=sk-...` + the key is authenticated and the catalog also holds the disabled plugins granted + to it through `object_permission.skills` on the key or its team. + Returns: Marketplace catalog with list of available plugins and their git sources. Example: ```bash claude plugin marketplace add http://localhost:4000/claude-code/marketplace.json + claude plugin marketplace add "http://localhost:4000/claude-code/marketplace.json?key=sk-..." claude plugin install my-plugin@litellm ``` """ try: prisma_client: Final = await _get_prisma_client() + caller: Final[UserAPIKeyAuth | None] = ( + await user_api_key_auth(request=request, api_key=f"Bearer {key}") if key else None + ) + visibility: Final[SkillVisibility] = skill_visibility(caller) plugins: Final[Sequence[_PluginRecord]] = await ClaudeCodePluginRepository(prisma_client).table.find_many( - where={"enabled": True} + where=visibility.where() ) plugin_list: Final = [] for plugin in plugins: + if not visibility.allows(plugin): + continue try: manifest: Mapping[str, object] = json.loads(plugin.manifest_json or "{}") except json.JSONDecodeError: @@ -149,7 +164,7 @@ async def get_marketplace(): return JSONResponse(content=marketplace) - except HTTPException: + except (HTTPException, ProxyException): raise except Exception as e: verbose_proxy_logger.exception("Error generating marketplace: %s", e) @@ -395,13 +410,15 @@ async def list_plugins( try: prisma_client: Final = await _get_prisma_client() - where: Final = {"enabled": True} if enabled_only else {} + visibility: Final[SkillVisibility] = skill_visibility(user_api_key_dict) plugins: Final[Sequence[_PluginRecord]] = await ClaudeCodePluginRepository(prisma_client).table.find_many( - where=where + where={"enabled": True} if enabled_only else visibility.where() ) plugin_list: Final = [] for p in plugins: + if not visibility.allows(p): + continue # Parse manifest to get additional fields manifest = json.loads(p.manifest_json) if p.manifest_json else {} @@ -473,6 +490,12 @@ async def get_plugin( detail={"error": f"Plugin '{plugin_name}' not found"}, ) + if not skill_visibility(user_api_key_dict).allows(plugin): + raise HTTPException( + status_code=403, + detail={"error": f"Plugin '{plugin_name}' is not granted to this key"}, + ) + manifest: Final[Mapping[str, object]] = json.loads(plugin.manifest_json or "{}") if plugin.manifest_json else {} return { diff --git a/litellm/proxy/anthropic_endpoints/claude_code_endpoints/claude_code_skill_access.py b/litellm/proxy/anthropic_endpoints/claude_code_endpoints/claude_code_skill_access.py new file mode 100644 index 00000000000..e1e1ee6f160 --- /dev/null +++ b/litellm/proxy/anthropic_endpoints/claude_code_endpoints/claude_code_skill_access.py @@ -0,0 +1,67 @@ +""" +Claude Code marketplace visibility: enabled plugins are public, disabled plugins +are private and resolve only for proxy admins or keys granted them via +``object_permission.skills``. +""" + +from dataclasses import dataclass +from typing import TYPE_CHECKING, Final, Protocol + +from litellm.proxy._types import LiteLLM_ObjectPermissionTable, UserAPIKeyAuth +from litellm.proxy.common_utils.resource_ownership import is_proxy_admin + +if TYPE_CHECKING: + from prisma.types import LiteLLM_ClaudeCodePluginTableWhereInput + + +class _SkillRecord(Protocol): + name: str + enabled: bool + + +def _skills_of(permission: LiteLLM_ObjectPermissionTable | None) -> frozenset[str]: + return frozenset(permission.skills or ()) if permission is not None else frozenset() + + +def granted_skills(user_api_key_dict: UserAPIKeyAuth) -> frozenset[str]: + """Key grant intersected with the team grant when both are non-empty; either alone applies as is. + + An empty list is the Prisma column default for every object-permission row, so it means + "no private grants configured here" and defers to the other scope, same as the agents check. + """ + key_skills: Final = _skills_of(user_api_key_dict.object_permission) + team_skills: Final = _skills_of(user_api_key_dict.team_object_permission) + match (bool(key_skills), bool(team_skills)): + case (True, True): + return key_skills & team_skills + case (True, False): + return key_skills + case _: + return team_skills + + +@dataclass(frozen=True, slots=True) +class SkillVisibility: + granted: frozenset[str] + sees_private: bool + + def allows(self, skill: _SkillRecord) -> bool: + return skill.enabled or self.sees_private or skill.name in self.granted + + def where(self) -> "LiteLLM_ClaudeCodePluginTableWhereInput": + if self.sees_private: + return {} + if not self.granted: + return {"enabled": True} + return {"OR": [{"enabled": True}, {"name": {"in": sorted(self.granted)}}]} + + +PUBLIC_ONLY: Final = SkillVisibility(granted=frozenset(), sees_private=False) + + +def skill_visibility(user_api_key_dict: UserAPIKeyAuth | None) -> SkillVisibility: + if user_api_key_dict is None: + return PUBLIC_ONLY + if is_proxy_admin(user_api_key_dict): + return SkillVisibility(granted=frozenset(), sees_private=True) + return SkillVisibility(granted=granted_skills(user_api_key_dict), sees_private=False) diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 05c5aad9303..817df082d8c 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -282,6 +282,7 @@ model LiteLLM_ObjectPermissionTable { mcp_toolsets String[] @default([]) // Toolset IDs granted to this key/team/user search_tools String[] @default([]) // search_tool_name values this key/team/user may call mcp_tool_search_enabled Boolean? + skills String[] @default([]) // Claude Code plugin names granted to this key/team beyond the public (enabled) set teams LiteLLM_TeamTable[] projects LiteLLM_ProjectTable[] verification_tokens LiteLLM_VerificationToken[] diff --git a/litellm/repositories/object_permission_repository.py b/litellm/repositories/object_permission_repository.py index 6b1f9c68e47..7736939c696 100644 --- a/litellm/repositories/object_permission_repository.py +++ b/litellm/repositories/object_permission_repository.py @@ -40,6 +40,7 @@ class ObjectPermissionRepository(BaseRepository[LiteLLM_ObjectPermissionTable]): blocked_tools: list[str] | None = None, mcp_toolsets: list[str] | None = None, search_tools: list[str] | None = None, + skills: list[str] | None = None, ) -> LiteLLM_ObjectPermissionTable: """Create a new object permission record.""" data: Final[dict[str, Any]] = {} @@ -63,6 +64,8 @@ class ObjectPermissionRepository(BaseRepository[LiteLLM_ObjectPermissionTable]): data["mcp_toolsets"] = mcp_toolsets if search_tools is not None: data["search_tools"] = search_tools + if skills is not None: + data["skills"] = skills return await self.create(data) @@ -79,6 +82,7 @@ class ObjectPermissionRepository(BaseRepository[LiteLLM_ObjectPermissionTable]): blocked_tools: list[str] | None = None, mcp_toolsets: list[str] | None = None, search_tools: list[str] | None = None, + skills: list[str] | None = None, ) -> LiteLLM_ObjectPermissionTable | None: """Update an object permission record.""" data: Final[dict[str, Any]] = {} @@ -102,6 +106,8 @@ class ObjectPermissionRepository(BaseRepository[LiteLLM_ObjectPermissionTable]): data["mcp_toolsets"] = mcp_toolsets if search_tools is not None: data["search_tools"] = search_tools + if skills is not None: + data["skills"] = skills return await self.update(object_permission_id, data, id_field="object_permission_id") diff --git a/litellm/types/object_permission.py b/litellm/types/object_permission.py index 1b391a3a1ef..68661aed891 100644 --- a/litellm/types/object_permission.py +++ b/litellm/types/object_permission.py @@ -8,7 +8,7 @@ can adopt the type without violating the SDK-must-not-import-from-proxy layering rule. """ -from typing_extensions import TypedDict +from typing_extensions import ReadOnly, TypedDict class ObjectPermissionDict(TypedDict, total=False): @@ -23,3 +23,4 @@ class ObjectPermissionDict(TypedDict, total=False): models: list[str] | None search_tools: list[str] | None mcp_tool_search_enabled: bool | None + skills: ReadOnly[list[str] | None] diff --git a/schema.prisma b/schema.prisma index 05c5aad9303..817df082d8c 100644 --- a/schema.prisma +++ b/schema.prisma @@ -282,6 +282,7 @@ model LiteLLM_ObjectPermissionTable { mcp_toolsets String[] @default([]) // Toolset IDs granted to this key/team/user search_tools String[] @default([]) // search_tool_name values this key/team/user may call mcp_tool_search_enabled Boolean? + skills String[] @default([]) // Claude Code plugin names granted to this key/team beyond the public (enabled) set teams LiteLLM_TeamTable[] projects LiteLLM_ProjectTable[] verification_tokens LiteLLM_VerificationToken[] diff --git a/tests/pass_through_unit_tests/test_claude_code_marketplace.py b/tests/pass_through_unit_tests/test_claude_code_marketplace.py index 2ca81f1d5d3..bedb8830559 100644 --- a/tests/pass_through_unit_tests/test_claude_code_marketplace.py +++ b/tests/pass_through_unit_tests/test_claude_code_marketplace.py @@ -216,7 +216,7 @@ async def test_get_marketplace(mock_prisma_client): ) # Now get the marketplace - response = await get_marketplace() + response = await get_marketplace(request=MagicMock()) # Response is a JSONResponse, get the body body = json.loads(response.body.decode()) diff --git a/tests/test_litellm/proxy/anthropic_endpoints/test_claude_code_marketplace.py b/tests/test_litellm/proxy/anthropic_endpoints/test_claude_code_marketplace.py index aa2da262bae..a7c2bd7ba20 100644 --- a/tests/test_litellm/proxy/anthropic_endpoints/test_claude_code_marketplace.py +++ b/tests/test_litellm/proxy/anthropic_endpoints/test_claude_code_marketplace.py @@ -11,17 +11,20 @@ from fastapi import HTTPException from unittest.mock import AsyncMock, MagicMock import litellm -from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy._types import LiteLLM_ObjectPermissionTable, ProxyException, UserAPIKeyAuth from litellm.proxy.proxy_server import LitellmUserRoles from litellm.types.proxy.claude_code_endpoints import ( RegisterPluginRequest, UpdatePluginRequest, ) +from litellm.proxy.anthropic_endpoints.claude_code_endpoints import claude_code_marketplace from litellm.proxy.anthropic_endpoints.claude_code_endpoints.claude_code_marketplace import ( delete_plugin, disable_plugin, enable_plugin, get_marketplace, + get_plugin, + list_plugins, register_plugin, update_plugin, ) @@ -38,11 +41,15 @@ def _make_mock_prisma(): async def _find_unique(where): return store.get(where.get("name")) + def _matches(record, where) -> bool: + if "OR" in where: + return any(_matches(record, clause) for clause in where["OR"]) + if "enabled" in where and record.enabled != where["enabled"]: + return False + return "name" not in where or record.name in where["name"]["in"] + async def _find_many(where=None): - records = list(store.values()) - if where and "enabled" in where: - return [r for r in records if r.enabled == where["enabled"]] - return records + return [r for r in store.values() if _matches(r, where or {})] async def _create(data): record = MagicMock() @@ -52,6 +59,8 @@ def _make_mock_prisma(): record.description = data.get("description") record.manifest_json = data.get("manifest_json", "{}") record.enabled = data.get("enabled", True) + record.created_at = data.get("created_at") + record.updated_at = data.get("updated_at") store[data["name"]] = record return record @@ -271,13 +280,89 @@ async def test_get_marketplace_skips_plugin_with_null_manifest(): table = litellm.proxy.proxy_server.prisma_client.db.litellm_claudecodeplugintable await table.create(data={"name": "null-manifest-plugin", "manifest_json": None, "enabled": True}) - response = await get_marketplace() + response = await get_marketplace(request=MagicMock()) assert response.status_code == 200 body = json.loads(response.body) assert [plugin["name"] for plugin in body["plugins"]] == ["good-plugin"] +def _granted_user(skills: list[str]) -> UserAPIKeyAuth: + return UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + api_key="sk-granted", + user_id="granted-user", + object_permission=LiteLLM_ObjectPermissionTable(object_permission_id="perm-1", skills=skills), + ) + + +async def _register_public_and_private_plugins() -> None: + for name, enabled in (("public-skill", True), ("private-skill", False)): + await register_plugin( + request=RegisterPluginRequest(name=name, source=_GIT_SUBDIR_SOURCE, version="1.0.0"), + user_api_key_dict=_USER, + ) + if not enabled: + await disable_plugin(plugin_name=name, user_api_key_dict=_USER) + + +async def _listed_names(user: UserAPIKeyAuth) -> set[str]: + response = await list_plugins(user_api_key_dict=user) + return {plugin.name for plugin in response.plugins} + + +@pytest.mark.asyncio +async def test_list_plugins_shows_disabled_plugin_only_to_granted_key_or_admin(): + await _register_public_and_private_plugins() + + assert await _listed_names(_NON_ADMIN_USER) == {"public-skill"} + assert await _listed_names(_granted_user(["other-skill"])) == {"public-skill"} + assert await _listed_names(_granted_user(["private-skill"])) == {"public-skill", "private-skill"} + assert await _listed_names(_USER) == {"public-skill", "private-skill"} + + +@pytest.mark.asyncio +async def test_get_plugin_returns_403_for_disabled_plugin_the_key_is_not_granted(): + await _register_public_and_private_plugins() + + with pytest.raises(HTTPException) as exc_info: + await get_plugin(plugin_name="private-skill", user_api_key_dict=_NON_ADMIN_USER) + assert exc_info.value.status_code == 403 + + assert (await get_plugin(plugin_name="public-skill", user_api_key_dict=_NON_ADMIN_USER))["name"] == "public-skill" + granted = await get_plugin(plugin_name="private-skill", user_api_key_dict=_granted_user(["private-skill"])) + assert granted["name"] == "private-skill" + assert granted["enabled"] is False + + +async def _marketplace_names(key: str | None) -> list[str]: + response = await get_marketplace(request=MagicMock(), key=key) + assert response.status_code == 200 + return sorted(plugin["name"] for plugin in json.loads(response.body)["plugins"]) + + +@pytest.mark.asyncio +async def test_get_marketplace_key_query_param_adds_granted_disabled_plugins(monkeypatch): + await _register_public_and_private_plugins() + keys = {"sk-granted": _granted_user(["private-skill"]), "sk-plain": _NON_ADMIN_USER} + + async def _fake_auth(request, api_key: str) -> UserAPIKeyAuth: + token = api_key.removeprefix("Bearer ") + if token not in keys: + raise ProxyException(message="invalid key", type="auth_error", param="key", code=401) + return keys[token] + + monkeypatch.setattr(claude_code_marketplace, "user_api_key_auth", _fake_auth) + + assert await _marketplace_names(None) == ["public-skill"] + assert await _marketplace_names("sk-plain") == ["public-skill"] + assert await _marketplace_names("sk-granted") == ["private-skill", "public-skill"] + + with pytest.raises(ProxyException) as exc_info: + await get_marketplace(request=MagicMock(), key="sk-bogus") + assert exc_info.value.code == "401" + + @pytest.mark.asyncio async def test_register_plugin_git_subdir_missing_url(): """git-subdir without url field raises HTTP 400.""" diff --git a/tests/test_litellm/proxy/anthropic_endpoints/test_claude_code_skill_access.py b/tests/test_litellm/proxy/anthropic_endpoints/test_claude_code_skill_access.py new file mode 100644 index 00000000000..dc4016d6db0 --- /dev/null +++ b/tests/test_litellm/proxy/anthropic_endpoints/test_claude_code_skill_access.py @@ -0,0 +1,79 @@ +"""Unit tests for claude_code_skill_access.py: key/team grant resolution.""" + +from unittest.mock import MagicMock + +import pytest + +from litellm.proxy._types import LiteLLM_ObjectPermissionTable, LitellmUserRoles, UserAPIKeyAuth +from litellm.proxy.anthropic_endpoints.claude_code_endpoints.claude_code_skill_access import ( + granted_skills, + skill_visibility, +) + + +def _perm(skills: list[str] | None) -> LiteLLM_ObjectPermissionTable: + return LiteLLM_ObjectPermissionTable(object_permission_id="perm", skills=skills) + + +def _key(key_skills: list[str] | None, team_skills: list[str] | None) -> UserAPIKeyAuth: + return UserAPIKeyAuth( + api_key="sk-user", + user_role=LitellmUserRoles.INTERNAL_USER, + object_permission=_perm(key_skills) if key_skills is not None else None, + team_object_permission=_perm(team_skills) if team_skills is not None else None, + ) + + +def _plugin(name: str, enabled: bool) -> MagicMock: + record = MagicMock() + record.name = name + record.enabled = enabled + return record + + +@pytest.mark.parametrize( + ("key_skills", "team_skills", "expected"), + [ + (None, None, frozenset()), + ([], [], frozenset()), + (["a", "b"], None, frozenset({"a", "b"})), + (None, ["a", "b"], frozenset({"a", "b"})), + (["a", "b"], ["b", "c"], frozenset({"b"})), + (["a"], ["c"], frozenset()), + (["a", "b"], [], frozenset({"a", "b"})), + ([], ["a", "b"], frozenset({"a", "b"})), + ], +) +def test_granted_skills_intersects_key_with_team(key_skills, team_skills, expected): + assert granted_skills(_key(key_skills, team_skills)) == expected + + +def test_visibility_enabled_plugin_is_public_for_everyone(): + public = _plugin("public-skill", enabled=True) + + assert skill_visibility(None).allows(public) + assert skill_visibility(_key(None, None)).allows(public) + assert skill_visibility(_key([], ["other"])).allows(public) + + +def test_visibility_disabled_plugin_needs_grant_or_admin(): + private = _plugin("private-skill", enabled=False) + admin = UserAPIKeyAuth(api_key="sk-1234", user_role=LitellmUserRoles.PROXY_ADMIN) + + assert not skill_visibility(None).allows(private) + assert not skill_visibility(_key(None, None)).allows(private) + assert not skill_visibility(_key(["other-skill"], None)).allows(private) + assert not skill_visibility(_key(["private-skill"], ["other-skill"])).allows(private) + assert skill_visibility(_key(["private-skill"], None)).allows(private) + assert skill_visibility(_key(None, ["private-skill"])).allows(private) + assert skill_visibility(admin).allows(private) + + +def test_where_clause_bounds_the_plugin_query_to_what_the_caller_may_see(): + admin = UserAPIKeyAuth(api_key="sk-1234", user_role=LitellmUserRoles.PROXY_ADMIN) + + assert skill_visibility(None).where() == {"enabled": True} + assert skill_visibility(_key(None, None)).where() == {"enabled": True} + assert skill_visibility(_key(["b", "a"], None)).where() == {"OR": [{"enabled": True}, {"name": {"in": ["a", "b"]}}]} + assert skill_visibility(_key(["a", "b"], ["b"])).where() == {"OR": [{"enabled": True}, {"name": {"in": ["b"]}}]} + assert skill_visibility(admin).where() == {} diff --git a/tests/test_litellm/proxy/auth/test_route_checks.py b/tests/test_litellm/proxy/auth/test_route_checks.py index 0fd19f518d9..4f2ff1582dd 100644 --- a/tests/test_litellm/proxy/auth/test_route_checks.py +++ b/tests/test_litellm/proxy/auth/test_route_checks.py @@ -3826,3 +3826,17 @@ def test_team_disable_logging_stays_proxy_admin_only(): def test_neighbouring_team_routes_stay_closed(route): """The grant is the callback paths and nothing else on the team namespace.""" assert "Only proxy admin" in _gate(route, LitellmUserRoles.INTERNAL_USER.value) + + +@pytest.mark.parametrize( + "route", + [ + "/claude-code/marketplace.json", + "/claude-code/plugins", + "/claude-code/plugins/my-skill", + ], +) +def test_claude_code_marketplace_routes_open_to_internal_users(route): + """Per-skill visibility is enforced inside the handler, so the route gate must let non-admins through.""" + assert RouteChecks.is_llm_api_route(route) is True + assert _gate(route, LitellmUserRoles.INTERNAL_USER.value) == "allowed" diff --git a/tests/test_litellm/proxy/management_endpoints/test_customer_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_customer_endpoints.py index 1225cb80224..bd59a82cbd2 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_customer_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_customer_endpoints.py @@ -763,6 +763,7 @@ _EXPECTED_CUSTOMER = { "blocked_tools": [], "search_tools": [], "mcp_tool_search_enabled": None, + "skills": None, }, } diff --git a/tests/test_litellm/proxy/management_helpers/test_object_permission_utils.py b/tests/test_litellm/proxy/management_helpers/test_object_permission_utils.py index d7ebb1f60bf..078315c2bf8 100644 --- a/tests/test_litellm/proxy/management_helpers/test_object_permission_utils.py +++ b/tests/test_litellm/proxy/management_helpers/test_object_permission_utils.py @@ -122,6 +122,29 @@ async def test_set_object_permission_persists_mcp_tool_search_enabled(): assert created_data["mcp_tool_search_enabled"] is True +@pytest.mark.asyncio +async def test_set_object_permission_persists_skills(): + mock_prisma_client = MagicMock() + mock_created_permission = MagicMock() + mock_created_permission.object_permission_id = "perm_id" + mock_prisma_client.db.litellm_objectpermissiontable.create = AsyncMock( + return_value=mock_created_permission + ) + + data_json = { + "object_permission": LiteLLM_ObjectPermissionBase(skills=["private-skill"]).model_dump(), + } + + await _set_object_permission(data_json=data_json, prisma_client=mock_prisma_client) + + created_data = ( + mock_prisma_client.db.litellm_objectpermissiontable.create.call_args.kwargs[ + "data" + ] + ) + assert created_data["skills"] == ["private-skill"] + + # ---- Tests for _extract_requested_mcp_server_ids ---- diff --git a/ui/litellm-dashboard/src/components/Teams.test.tsx b/ui/litellm-dashboard/src/components/Teams.test.tsx index c7df35197e6..70c30596c75 100644 --- a/ui/litellm-dashboard/src/components/Teams.test.tsx +++ b/ui/litellm-dashboard/src/components/Teams.test.tsx @@ -33,6 +33,14 @@ vi.mock("./mcp_server_management/MCPServerSelector", () => ({ ), })); +vi.mock("./skills/SkillSelector", () => ({ + default: ({ onChange }: { onChange: (selected: string[]) => void }) => ( + + ), +})); + const can = vi.fn(); vi.mock("@/app/(dashboard)/hooks/useCan", () => ({ default: (...args: unknown[]) => can(...args), @@ -1359,6 +1367,27 @@ describe("Teams - the exact bytes the create call sends", () => { }); }); + it("puts the selected skills into object_permission.skills and drops the form key", async () => { + await openCreateModal(); + await openSection("Skill Settings", /Allowed Skills/); + fireEvent.click(screen.getByTestId("select-private-skill")); + + const payload = await submit(); + + expect(payload.object_permission).toStrictEqual({ skills: ["private-skill"] }); + expect(payload).not.toHaveProperty("object_permission_skills"); + }); + + it("sends no object_permission when Skill Settings is opened but nothing is selected", async () => { + await openCreateModal(); + await openSection("Skill Settings", /Allowed Skills/); + + const payload = await submit(); + + expect(payload).not.toHaveProperty("object_permission"); + expect(payload).not.toHaveProperty("object_permission_skills"); + }); + it("includes selected MCP toolsets in the create object permission", async () => { await openCreateModal(); await openSection("MCP Settings", /Allowed MCP Servers/); diff --git a/ui/litellm-dashboard/src/components/Teams.tsx b/ui/litellm-dashboard/src/components/Teams.tsx index c4060163c78..d2d18139893 100644 --- a/ui/litellm-dashboard/src/components/Teams.tsx +++ b/ui/litellm-dashboard/src/components/Teams.tsx @@ -50,6 +50,7 @@ import { Organization, getDefaultTeamSettings, getGuardrailsList, getPoliciesLis import NumericalInput from "./shared/numerical_input"; import VectorStoreSelector from "./vector_store_management/VectorStoreSelector"; import SearchToolSelector from "./search_tools/SearchToolSelector"; +import SkillSelector from "./skills/SkillSelector"; import { Dialog, DialogContent, DialogHeader, DialogTitle } from "@/components/ui/dialog"; interface TeamProps { @@ -99,6 +100,7 @@ const teamCreateFieldsSchema = z.object({ mcp_tool_permissions: z.record(z.string(), z.array(z.string())).optional(), allowed_agents_and_groups: z.object({ agents: z.array(z.string()), accessGroups: z.array(z.string()) }).optional(), object_permission_search_tools: z.array(z.string()).optional(), + object_permission_skills: z.array(z.string()).optional(), }); type TeamCreateFormValues = z.infer; @@ -128,6 +130,7 @@ const EMPTY_TEAM_CREATE_VALUES: TeamCreateFormValues = { mcp_tool_permissions: {}, allowed_agents_and_groups: undefined, object_permission_search_tools: undefined, + object_permission_skills: undefined, }; const ADDITIONAL_SETTINGS_FIELDS = [ @@ -147,6 +150,7 @@ const ADDITIONAL_SETTINGS_FIELDS = [ const MCP_SETTINGS_FIELDS = ["allowed_mcp_servers_and_groups", "mcp_tool_permissions"] as const; const AGENT_SETTINGS_FIELDS = ["allowed_agents_and_groups"] as const; const SEARCH_TOOL_SETTINGS_FIELDS = ["object_permission_search_tools"] as const; +const SKILL_SETTINGS_FIELDS = ["object_permission_skills"] as const; const isParsableJson = (value: string | undefined): boolean => { if (!value) { @@ -214,6 +218,7 @@ const Teams: React.FC = ({ accessToken, userID, userRole, premiumUser const [mcpSettingsOpen, setMcpSettingsOpen] = useState(false); const [agentSettingsOpen, setAgentSettingsOpen] = useState(false); const [searchToolSettingsOpen, setSearchToolSettingsOpen] = useState(false); + const [skillSettingsOpen, setSkillSettingsOpen] = useState(false); const adminOrgs = useMemo( () => getAdminOrganizations(userRole, userID, organizations), @@ -505,6 +510,14 @@ const Teams: React.FC = ({ accessToken, userID, userRole, premiumUser delete formValues.object_permission_search_tools; } + if (Array.isArray(formValues.object_permission_skills) && formValues.object_permission_skills.length > 0) { + if (!formValues.object_permission) { + formValues.object_permission = {}; + } + formValues.object_permission.skills = formValues.object_permission_skills; + } + delete formValues.object_permission_skills; + // Add model_aliases if any are defined if (Object.keys(modelAliases).length > 0) { formValues.model_aliases = modelAliases; @@ -540,6 +553,7 @@ const Teams: React.FC = ({ accessToken, userID, userRole, premiumUser ...(mcpSettingsOpen ? [] : MCP_SETTINGS_FIELDS), ...(agentSettingsOpen ? [] : AGENT_SETTINGS_FIELDS), ...(searchToolSettingsOpen ? [] : SEARCH_TOOL_SETTINGS_FIELDS), + ...(skillSettingsOpen ? [] : SKILL_SETTINGS_FIELDS), ]); return Object.fromEntries(Object.entries(values).filter(([key]) => !unmounted.has(key))); }; @@ -1163,6 +1177,38 @@ const Teams: React.FC = ({ accessToken, userID, userRole, premiumUser + + + Skill Settings + + + + + {({ value, onChange }) => ( + + )} + + + + Logging Settings diff --git a/ui/litellm-dashboard/src/components/object_permissions_view.tsx b/ui/litellm-dashboard/src/components/object_permissions_view.tsx index e90645a9e9f..ae46729222a 100644 --- a/ui/litellm-dashboard/src/components/object_permissions_view.tsx +++ b/ui/litellm-dashboard/src/components/object_permissions_view.tsx @@ -30,6 +30,7 @@ export function ObjectPermissionsView({ const agents = objectPermission?.agents || []; const agentAccessGroups = objectPermission?.agent_access_groups || []; const searchTools = objectPermission?.search_tools || []; + const skills = objectPermission?.skills || []; const content = (
@@ -58,6 +59,16 @@ export function ObjectPermissionsView({

{searchTools.join(", ")}

)}
+
+

Skills

+ {skills.length === 0 ? ( +

+ No private skills granted. Only enabled (public) Claude Code plugins are visible. +

+ ) : ( +

{skills.join(", ")}

+ )} +
); diff --git a/ui/litellm-dashboard/src/components/organisms/createKeyPayload.test.ts b/ui/litellm-dashboard/src/components/organisms/createKeyPayload.test.ts index 8b7283370a2..fef94cc3c2b 100644 --- a/ui/litellm-dashboard/src/components/organisms/createKeyPayload.test.ts +++ b/ui/litellm-dashboard/src/components/organisms/createKeyPayload.test.ts @@ -301,6 +301,16 @@ describe("object_permission", () => { ).toStrictEqual(aliasOnly({ object_permission: { agents: ["a-1"], agent_access_groups: ["ag-1"] } })); }); + it("moves the selected skills under object_permission and off the top level", () => { + expect(payloadOf(build({ key_alias: "my-key", allowed_skills: ["private-skill"] }))).toStrictEqual( + aliasOnly({ object_permission: { skills: ["private-skill"] } }), + ); + }); + + it("sends no object_permission for an empty skill selection", () => { + expect(payloadOf(build({ key_alias: "my-key", allowed_skills: [] }))).toStrictEqual(aliasOnly()); + }); + it("merges every source into a single object_permission", () => { const everySource = { key_alias: "my-key", @@ -308,6 +318,7 @@ describe("object_permission", () => { allowed_mcp_servers_and_groups: { servers: ["s-1"], accessGroups: ["g-1"], toolsets: ["t-1"] }, mcp_tool_permissions: { "s-1": ["read"] }, allowed_agents_and_groups: { agents: ["a-1"], accessGroups: ["ag-1"] }, + allowed_skills: ["private-skill"], }; expect(payloadOf(build(everySource))).toStrictEqual( aliasOnly({ @@ -319,6 +330,7 @@ describe("object_permission", () => { mcp_tool_permissions: { "s-1": ["read"] }, agents: ["a-1"], agent_access_groups: ["ag-1"], + skills: ["private-skill"], }, }), ); diff --git a/ui/litellm-dashboard/src/components/organisms/createKeyPayload.ts b/ui/litellm-dashboard/src/components/organisms/createKeyPayload.ts index a528dbaeb2c..37a973d5def 100644 --- a/ui/litellm-dashboard/src/components/organisms/createKeyPayload.ts +++ b/ui/litellm-dashboard/src/components/organisms/createKeyPayload.ts @@ -112,6 +112,7 @@ interface PermissionSources { readonly toolPermissions: unknown | undefined; readonly extraMcpAccessGroups: unknown[] | undefined; readonly agents: AgentSelection | undefined; + readonly skills: unknown[] | undefined; } const readPermissionSources = (values: Record): PermissionSources => ({ @@ -120,6 +121,7 @@ const readPermissionSources = (values: Record): PermissionSourc toolPermissions: readToolPermissions(values.mcp_tool_permissions), extraMcpAccessGroups: nonEmptyList(values.allowed_mcp_access_groups), agents: readAgentSelection(values.allowed_agents_and_groups), + skills: nonEmptyList(values.allowed_skills), }); const buildObjectPermission = ({ @@ -128,6 +130,7 @@ const buildObjectPermission = ({ toolPermissions, extraMcpAccessGroups, agents, + skills, }: PermissionSources): Record | undefined => { const permission: Record = { ...(vectorStores && { vector_stores: vectorStores }), @@ -138,6 +141,7 @@ const buildObjectPermission = ({ ...(extraMcpAccessGroups && { mcp_access_groups: extraMcpAccessGroups }), ...(agents?.agents && { agents: agents.agents }), ...(agents?.accessGroups && { agent_access_groups: agents.accessGroups }), + ...(skills && { skills }), }; return Object.keys(permission).length > 0 ? permission : undefined; }; @@ -148,6 +152,7 @@ const consumedSourceKeys = ( ): ReadonlySet => new Set([ "mcp_tool_permissions", + "allowed_skills", ...(values.disable_global_guardrails ? [] : ["disable_global_guardrails"]), ...(vectorStores ? ["allowed_vector_store_ids"] : []), ...(mcp ? ["allowed_mcp_servers_and_groups"] : []), diff --git a/ui/litellm-dashboard/src/components/organisms/create_key_button.integration.test.tsx b/ui/litellm-dashboard/src/components/organisms/create_key_button.integration.test.tsx index 3471ef00eb0..2bf39bf4cd7 100644 --- a/ui/litellm-dashboard/src/components/organisms/create_key_button.integration.test.tsx +++ b/ui/litellm-dashboard/src/components/organisms/create_key_button.integration.test.tsx @@ -84,6 +84,13 @@ vi.mock("../networking", async (importOriginal) => { getPossibleUserRoles: vi.fn().mockResolvedValue({}), userFilterUICall: vi.fn().mockResolvedValue([]), getAgentsList: vi.fn().mockResolvedValue({ agents: [] }), + getClaudeCodePluginsList: vi.fn().mockResolvedValue({ + plugins: [ + { name: "public-skill", enabled: true }, + { name: "private-skill", enabled: false }, + ], + count: 2, + }), getPassThroughEndpointsCall: vi.fn().mockResolvedValue({ endpoints: [] }), vectorStoreListCall: vi.fn().mockResolvedValue({ data: [] }), listMCPTools: vi.fn().mockResolvedValue(emptyMcpTools), @@ -109,6 +116,7 @@ const OPENAPI_SCHEMA = { const SECTIONS = { mcp: /MCP Settings/i, agent: /Agent Settings/i, + skill: /Skill Settings/i, logging: /Logging Settings/i, router: /Router Settings/i, aliases: /Model Aliases/i, @@ -164,6 +172,7 @@ const ROUTER_SETTINGS_DEFAULT = { const SECTION_PAYLOAD_ADDITIONS: Record> = { mcp: { allowed_mcp_servers_and_groups: { servers: [], accessGroups: [] } }, agent: { allowed_agents_and_groups: undefined }, + skill: {}, logging: {}, router: { router_settings: ROUTER_SETTINGS_DEFAULT }, aliases: {}, @@ -329,6 +338,21 @@ describe("CreateKey", () => { expect(Object.keys(serialised).sort()).toStrictEqual([...wireKeys].sort()); }); + it("moves a picked private skill under object_permission.skills and off the top level", async () => { + await openModal(); + await nameTheKey(); + await openSection(/Optional Settings/i); + await openSection(SECTIONS.skill); + await userEvent.click(await screen.findByRole("combobox", { name: "Select skills (optional)" })); + await userEvent.click(await screen.findByRole("option", { name: "private-skill (private)" })); + await userEvent.keyboard("{Escape}"); + await submit(); + + const payload = await createdPayload(); + expect(payload.object_permission).toStrictEqual({ skills: ["private-skill"] }); + expect(payload).not.toHaveProperty("allowed_skills"); + }); + it("omits a budget typed into a section the user closed again, rather than sending it as null", async () => { await openModal(); await nameTheKey(); diff --git a/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx b/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx index 16ea99ea94c..831cb0cf6a6 100644 --- a/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx +++ b/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx @@ -27,6 +27,7 @@ import React, { useEffect, useMemo, useRef, useState } from "react"; import { type Control, useForm, useWatch, type UseFormSetValue } from "react-hook-form"; import { rolesWithWriteAccess } from "../../utils/roles"; import AgentSelector from "../agent_management/AgentSelector"; +import SkillSelector from "../skills/SkillSelector"; import AccessGroupSelector from "../common_components/AccessGroupSelector"; import BudgetDurationDropdown from "../common_components/budget_duration_dropdown"; import SchemaFormFields from "../common_components/check_openapi_schema"; @@ -1557,6 +1558,36 @@ const CreateKey: React.FC = ({ team, teams, data, addKey, autoOp + + + Skill Settings + + + + + Allowed Skills{" "} + + + + + } + name="allowed_skills" + help="Select private skills this key can access in the Claude Code marketplace" + > + {(control) => ( + + )} + + + + {premiumUser ? ( diff --git a/ui/litellm-dashboard/src/components/skills/SkillSelector.test.tsx b/ui/litellm-dashboard/src/components/skills/SkillSelector.test.tsx new file mode 100644 index 00000000000..59aff7c25de --- /dev/null +++ b/ui/litellm-dashboard/src/components/skills/SkillSelector.test.tsx @@ -0,0 +1,54 @@ +import userEvent from "@testing-library/user-event"; +import { beforeEach, describe, expect, it, vi } from "vitest"; +import { renderWithProviders, screen } from "../../../tests/test-utils"; +import { getClaudeCodePluginsList } from "../networking"; +import SkillSelector from "./SkillSelector"; + +vi.mock("../networking", () => ({ + getClaudeCodePluginsList: vi.fn(), +})); + +describe("SkillSelector", () => { + beforeEach(() => { + vi.clearAllMocks(); + vi.mocked(getClaudeCodePluginsList).mockResolvedValue({ + plugins: [ + { name: "public-skill", enabled: true }, + { name: "private-skill", enabled: false }, + ], + count: 2, + }); + }); + + it("should list marketplace plugins and mark disabled ones as private", async () => { + const user = userEvent.setup(); + renderWithProviders(); + + await user.click(screen.getByRole("combobox")); + + expect(await screen.findByRole("option", { name: "public-skill" })).toBeInTheDocument(); + expect(screen.getByRole("option", { name: "private-skill (private)" })).toBeInTheDocument(); + expect(getClaudeCodePluginsList).toHaveBeenCalledWith("token"); + }); + + it("should report the selected skill names", async () => { + const user = userEvent.setup(); + const onChange = vi.fn(); + renderWithProviders(); + + await user.click(screen.getByRole("combobox")); + await user.click(await screen.findByRole("option", { name: "private-skill (private)" })); + + expect(onChange).toHaveBeenCalledWith(["private-skill"]); + }); + + it("should clear all selected skills", async () => { + const user = userEvent.setup(); + const onChange = vi.fn(); + renderWithProviders(); + + await user.click(screen.getByRole("button", { name: "Clear all skills" })); + + expect(onChange).toHaveBeenCalledWith([]); + }); +}); diff --git a/ui/litellm-dashboard/src/components/skills/SkillSelector.tsx b/ui/litellm-dashboard/src/components/skills/SkillSelector.tsx new file mode 100644 index 00000000000..b21c8b7e872 --- /dev/null +++ b/ui/litellm-dashboard/src/components/skills/SkillSelector.tsx @@ -0,0 +1,109 @@ +import React, { useEffect, useState } from "react"; +import { + Combobox, + ComboboxChip, + ComboboxChips, + ComboboxChipsInput, + ComboboxClear, + ComboboxContent, + ComboboxEmpty, + ComboboxItem, + ComboboxList, + ComboboxValue, + useComboboxAnchor, +} from "@/components/ui/combobox"; +import { cn } from "@/lib/cva.config"; +import { getClaudeCodePluginsList } from "../networking"; + +export interface SkillSelectorProps { + onChange: (selected: string[]) => void; + value?: string[]; + className?: string; + accessToken: string; + placeholder?: string; + disabled?: boolean; +} + +interface SkillOption { + readonly name: string; + readonly enabled: boolean; +} + +const readSkillOptions = (data: unknown): SkillOption[] => { + const plugins = (data as { plugins?: unknown[] } | undefined)?.plugins; + if (!Array.isArray(plugins)) return []; + return plugins.flatMap((plugin) => { + const record = plugin as { name?: unknown; enabled?: unknown }; + return typeof record.name === "string" && record.name.length > 0 + ? [{ name: record.name, enabled: record.enabled !== false }] + : []; + }); +}; + +const SkillSelector: React.FC = ({ + onChange, + value, + className, + accessToken, + placeholder = "Select skills (optional)", + disabled = false, +}) => { + const anchor = useComboboxAnchor(); + const [options, setOptions] = useState([]); + const [loading, setLoading] = useState(false); + + useEffect(() => { + const load = async () => { + if (!accessToken) return; + setLoading(true); + try { + setOptions(readSkillOptions(await getClaudeCodePluginsList(accessToken))); + } catch (e) { + console.error("Failed to load skills:", e); + } finally { + setLoading(false); + } + }; + load(); + }, [accessToken]); + + return ( + option.name)} + value={value ?? []} + onValueChange={(selected: string[]) => onChange(selected)} + disabled={disabled} + > + } className={cn("w-full", className)} aria-busy={loading}> + + {(selected: string[]) => + selected.map((skill) => ( + + {skill} + + )) + } + + + {value && value.length > 0 && } + + + {loading ? "Loading skills…" : "No skills found"} + + {(skill: string) => { + const isPrivate = options.some((option) => option.name === skill && !option.enabled); + return ( + + {skill} + {isPrivate && private} + + ); + }} + + + + ); +}; + +export default SkillSelector; diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx index b286c8c1303..a25ca28651a 100644 --- a/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx @@ -52,6 +52,7 @@ vi.mock("@/components/networking", () => ({ listMCPTools: vi.fn().mockResolvedValue({ tools: [] }), vectorStoreListCall: vi.fn().mockResolvedValue({ data: [] }), getAgentsList: vi.fn().mockResolvedValue({ agents: [] }), + getClaudeCodePluginsList: vi.fn().mockResolvedValue({ plugins: [], count: 0 }), })); const can = vi.fn(); @@ -1990,6 +1991,7 @@ describe("TeamInfoView - the exact bytes the update call sends", () => { agents: [], agent_access_groups: [], vector_stores: ["vs-1"], + skills: [], }; it("leaves every team member key out of the request body for an untouched save with both sections closed", async () => { @@ -2052,6 +2054,47 @@ describe("TeamInfoView - the exact bytes the update call sends", () => { expect(objectPermission.agent_access_groups).toStrictEqual([]); }); + it("resends the stored skills when the selector is left untouched", async () => { + const user = userEvent.setup({ delay: null }); + vi.mocked(networking.teamInfoCall).mockResolvedValue( + createMockTeamData({ models: ["gpt-4"], object_permission: { skills: ["private-skill"] } }), + ); + vi.mocked(networking.teamUpdateCall).mockResolvedValue({ data: {}, team_id: "123" } as any); + + renderWithProviders(); + await waitFor(() => expect(screen.queryAllByText("Test Team").length).toBeGreaterThan(0)); + await user.click(screen.getByRole("tab", { name: "Settings" })); + await user.click(await screen.findByRole("button", { name: /edit settings/i })); + await screen.findByLabelText("Team Name"); + + const payload = await save(user); + + const objectPermission = wireBody(payload).object_permission as Record; + expect(objectPermission.skills).toStrictEqual(["private-skill"]); + }); + + it("sends an empty skills array after the last skill chip is removed", async () => { + const user = userEvent.setup({ delay: null }); + vi.mocked(networking.teamInfoCall).mockResolvedValue( + createMockTeamData({ models: ["gpt-4"], object_permission: { skills: ["private-skill"] } }), + ); + vi.mocked(networking.teamUpdateCall).mockResolvedValue({ data: {}, team_id: "123" } as any); + + renderWithProviders(); + await waitFor(() => expect(screen.queryAllByText("Test Team").length).toBeGreaterThan(0)); + await user.click(screen.getByRole("tab", { name: "Settings" })); + await user.click(await screen.findByRole("button", { name: /edit settings/i })); + await screen.findByLabelText("Team Name"); + + await user.click(within(screen.getByLabelText("private-skill")).getByRole("button")); + expect(screen.queryByLabelText("private-skill")).not.toBeInTheDocument(); + + const payload = await save(user); + + const objectPermission = wireBody(payload).object_permission as Record; + expect(objectPermission.skills).toStrictEqual([]); + }); + it("sends an empty vector_stores array after the last vector store chip is removed", async () => { const user = userEvent.setup({ delay: null }); await openEditor(user); diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx index 4947cec1441..988e2e082aa 100644 --- a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx @@ -89,6 +89,7 @@ import ObjectPermissionsView from "../object_permissions_view"; import NumericalInput from "../shared/numerical_input"; import VectorStoreSelector from "../vector_store_management/VectorStoreSelector"; import SearchToolSelector from "../search_tools/SearchToolSelector"; +import SkillSelector from "../skills/SkillSelector"; import EditLoggingSettings from "./EditLoggingSettings"; import RouterSettingsAccordion, { RouterSettingsAccordionRef } from "../common_components/RouterSettingsAccordion"; import MemberModal from "./EditMembership"; @@ -371,6 +372,7 @@ const teamUpdateFieldsSchema = z.object({ mcp_tool_permissions: z.record(z.string(), z.array(z.string())).optional(), agents_and_groups: z.object({ agents: z.array(z.string()), accessGroups: z.array(z.string()) }).optional(), object_permission_search_tools: z.array(z.string()).optional(), + object_permission_skills: z.array(z.string()).optional(), organization_id: z.string().nullish(), logging_settings: z.array(z.unknown()).optional(), secret_manager_settings: z.string().optional(), @@ -419,6 +421,7 @@ const EMPTY_TEAM_UPDATE_VALUES: TeamUpdateFormValues = { mcp_tool_permissions: {}, agents_and_groups: { agents: [], accessGroups: [] }, object_permission_search_tools: [], + object_permission_skills: [], organization_id: null, logging_settings: [], secret_manager_settings: "", @@ -485,6 +488,7 @@ const toTeamFormValues = (info: TeamInfoRecord, effectiveGuardrails: string[]): accessGroups: info.object_permission?.agent_access_groups || [], }, object_permission_search_tools: info.object_permission?.search_tools || [], + object_permission_skills: info.object_permission?.skills || [], organization_id: info.organization_id, logging_settings: info.metadata?.logging || [], secret_manager_settings: info.metadata?.secret_manager_settings @@ -1047,6 +1051,10 @@ const TeamInfoView: React.FC = ({ updateData.object_permission.search_tools = values.object_permission_search_tools; } + if (Array.isArray(values.object_permission_skills)) { + updateData.object_permission.skills = values.object_permission_skills; + } + // Pass access_group_ids to the update request if (values.access_group_ids !== undefined) { updateData.access_group_ids = values.access_group_ids; @@ -1848,6 +1856,24 @@ const TeamInfoView: React.FC = ({ + + {({ value, onChange }) => ( + + )} + + {({ id, value, onChange }) => ( ( <> @@ -53,6 +55,36 @@ export const KeyTypeSelect = ({ ); +const SKILLS_HINT = + "Enabled skills are visible to every key. Grant disabled (private) Claude Code plugins to this key here."; + +export const KeyAgentAndSkillFields = ({ + control, + accessToken, +}: { + control: Control; + accessToken: string; +}) => ( + <> + + {({ value, onChange }) => ( + + )} + + + + {({ value, onChange }) => ( + + )} + + +); + export const KeyBudgetNumberField = ({ control, name, diff --git a/ui/litellm-dashboard/src/components/templates/KeyInfoView.handleKeyUpdate.test.tsx b/ui/litellm-dashboard/src/components/templates/KeyInfoView.handleKeyUpdate.test.tsx index 8a42ecff50c..522af5a85ad 100644 --- a/ui/litellm-dashboard/src/components/templates/KeyInfoView.handleKeyUpdate.test.tsx +++ b/ui/litellm-dashboard/src/components/templates/KeyInfoView.handleKeyUpdate.test.tsx @@ -366,6 +366,45 @@ describe("KeyInfoView handleKeyUpdate mcp_toolsets", () => { }); }); +describe("KeyInfoView handleKeyUpdate skills", () => { + it("should forward the skills the edit form supplies into object_permission and drop the form key", async () => { + renderView(true); + + fireEvent.click(screen.getByText("Settings")); + fireEvent.click(screen.getByText("Edit Settings")); + (globalThis as any).__TEST_FORM_VALUES = { + token: "tok_123", + skills: ["private-skill"], + }; + + fireEvent.click(screen.getByText("Mock Submit")); + + await waitFor(() => expect(keyUpdateCallMock).toHaveBeenCalled()); + + const [, sentPayload] = keyUpdateCallMock.mock.calls[0]; + expect(sentPayload.object_permission.skills).toEqual(["private-skill"]); + expect(sentPayload).not.toHaveProperty("skills"); + }); + + it("should send an explicit empty skills list when the form clears every skill", async () => { + renderView(true); + + fireEvent.click(screen.getByText("Settings")); + fireEvent.click(screen.getByText("Edit Settings")); + (globalThis as any).__TEST_FORM_VALUES = { + token: "tok_123", + skills: [], + }; + + fireEvent.click(screen.getByText("Mock Submit")); + + await waitFor(() => expect(keyUpdateCallMock).toHaveBeenCalled()); + + const [, sentPayload] = keyUpdateCallMock.mock.calls[0]; + expect(sentPayload.object_permission.skills).toEqual([]); + }); +}); + describe("KeyInfoView handleKeyUpdate budget_duration", () => { it("should send a canonical budget_duration through unchanged", async () => { renderView(true); diff --git a/ui/litellm-dashboard/src/components/templates/keyEditFormValues.ts b/ui/litellm-dashboard/src/components/templates/keyEditFormValues.ts index f24fb6e5a86..fec8e749143 100644 --- a/ui/litellm-dashboard/src/components/templates/keyEditFormValues.ts +++ b/ui/litellm-dashboard/src/components/templates/keyEditFormValues.ts @@ -46,6 +46,7 @@ export interface KeyEditFormValues { mcp_servers_and_groups?: McpServersAndGroups; mcp_tool_permissions?: Record; agents_and_groups?: AgentsAndGroups; + skills?: string[]; organization_id?: string | null; team_id?: string | null; logging_settings?: unknown[]; @@ -102,6 +103,7 @@ export const toKeyEditFormValues = (keyData: KeyResponse): KeyEditFormValues => agents: keyData.object_permission?.agents || [], accessGroups: keyData.object_permission?.agent_access_groups || [], }, + skills: keyData.object_permission?.skills || [], organization_id: keyData.organization_id, team_id: keyData.team_id, logging_settings: extractLoggingSettings(keyData.metadata), @@ -148,6 +150,7 @@ export const keyEditFormSchema = z.object({ mcp_servers_and_groups: z.custom(), mcp_tool_permissions: z.custom | undefined>(), agents_and_groups: z.custom(), + skills: z.custom(), organization_id: z.custom(), team_id: z.custom(), logging_settings: z.custom(), @@ -196,6 +199,7 @@ export const toSubmittedValues = ( mcp_servers_and_groups: values.mcp_servers_and_groups, mcp_tool_permissions: values.mcp_tool_permissions, agents_and_groups: values.agents_and_groups, + skills: values.skills, organization_id: values.organization_id, team_id: values.team_id, logging_settings: values.logging_settings, diff --git a/ui/litellm-dashboard/src/components/templates/key_edit_view.test.tsx b/ui/litellm-dashboard/src/components/templates/key_edit_view.test.tsx index 35e9268e1c2..b2c2a381b42 100644 --- a/ui/litellm-dashboard/src/components/templates/key_edit_view.test.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_edit_view.test.tsx @@ -59,6 +59,7 @@ vi.mock("../networking", async () => { agents: [], }), getAgentAccessGroups: vi.fn().mockResolvedValue([]), + getClaudeCodePluginsList: vi.fn().mockResolvedValue({ plugins: [], count: 0 }), }; }); @@ -135,6 +136,14 @@ vi.mock("../agent_management/AgentSelector", () => ({ ), })); +vi.mock("../skills/SkillSelector", () => ({ + default: ({ onChange }: { onChange: (selected: string[]) => void }) => ( + + ), +})); + vi.mock("../common_components/AccessGroupSelector", () => ({ default: ({ value = [], onChange }: { value?: string[]; onChange?: (v: string[]) => void }) => ( { mcp_servers_and_groups: { servers: [], accessGroups: [], toolsets: [] }, mcp_tool_permissions: {}, agents_and_groups: { agents: [], accessGroups: [] }, + skills: [], organization_id: null, team_id: null, logging_settings: [], @@ -2241,6 +2251,36 @@ describe("KeyEditView", () => { expect(onSubmitMock.mock.calls[0][0].agents_and_groups.agents).toEqual(["agent-1"]); }); + it("carries a picked skill into the payload", async () => { + const onSubmitMock = vi.fn().mockResolvedValue(undefined); + renderForPayload(onSubmitMock); + await screen.findByRole("button", { name: /save changes/i }); + + await userEvent.click(screen.getByRole("button", { name: "pick skill" })); + await userEvent.click(screen.getByRole("button", { name: /save changes/i })); + + await waitFor(() => { + expect(onSubmitMock).toHaveBeenCalled(); + }); + expect(onSubmitMock.mock.calls[0][0].skills).toEqual(["private-skill"]); + }); + + it("preloads the stored skills into the payload when the selector is left untouched", async () => { + const onSubmitMock = vi.fn().mockResolvedValue(undefined); + renderForPayload(onSubmitMock, { + ...MOCK_KEY_DATA, + object_permission: { ...MOCK_KEY_DATA.object_permission, skills: ["stored-skill"] }, + } as KeyResponse); + await screen.findByRole("button", { name: /save changes/i }); + + await userEvent.click(screen.getByRole("button", { name: /save changes/i })); + + await waitFor(() => { + expect(onSubmitMock).toHaveBeenCalled(); + }); + expect(onSubmitMock.mock.calls[0][0].skills).toEqual(["stored-skill"]); + }); + it("carries an added logging integration into the payload", async () => { const onSubmitMock = vi.fn().mockResolvedValue(undefined); renderForPayload(onSubmitMock); diff --git a/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx b/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx index b1bbbeed6f9..e464dfe5008 100644 --- a/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx @@ -15,7 +15,6 @@ import { FormField } from "@/components/shared/form/FormField"; import React, { useEffect, useRef, useState } from "react"; import { hasCapability } from "../../utils/capabilities"; import { isProxyAdminRole, rolesWithWriteAccess } from "../../utils/roles"; -import AgentSelector from "../agent_management/AgentSelector"; import AccessGroupSelector from "../common_components/AccessGroupSelector"; import BudgetDurationDropdown from "../common_components/budget_duration_dropdown"; import { mapInternalToDisplayNames } from "../callback_info_helpers"; @@ -32,9 +31,8 @@ import { modelSentinelOptions, parseAllowedRoutes, } from "./keyEditFieldNormalizers"; -import { KeyBudgetNumberField, KeyTypeSelect, labelWithHint } from "./KeyEditViewControls"; +import { KeyAgentAndSkillFields, KeyBudgetNumberField, KeyTypeSelect, labelWithHint } from "./KeyEditViewControls"; import { - AgentsAndGroups, KeyEditFormValues, keyEditFormSchema, McpServersAndGroups, @@ -762,16 +760,7 @@ export function KeyEditView({ /> - - {({ value, onChange }) => ( - - )} - + * - claude plugin install @ * + * Without `key` the catalog holds the enabled (public) plugins. With `?key=sk-...` + * the key is authenticated and the catalog also holds the disabled plugins granted + * to it through `object_permission.skills` on the key or its team. + * * Returns: * Marketplace catalog with list of available plugins and their git sources. * * Example: * ```bash * claude plugin marketplace add http://localhost:4000/claude-code/marketplace.json + * claude plugin marketplace add "http://localhost:4000/claude-code/marketplace.json?key=sk-..." * claude plugin install my-plugin@litellm * ``` */ @@ -29333,6 +29338,8 @@ export interface components { models?: string[] | null; /** Search Tools */ search_tools?: string[] | null; + /** Skills */ + skills?: string[] | null; /** Vector Stores */ vector_stores?: string[] | null; }; @@ -29386,6 +29393,8 @@ export interface components { * @default [] */ search_tools: string[] | null; + /** Skills */ + skills?: string[] | null; /** * Vector Stores * @default [] @@ -43206,7 +43215,9 @@ export interface operations { }; get_marketplace_claude_code_marketplace_json_get: { parameters: { - query?: never; + query?: { + key?: string | null; + }; header?: never; path?: never; cookie?: never; @@ -43222,6 +43233,15 @@ export interface operations { "application/json": unknown; }; }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; }; }; list_plugins_claude_code_plugins_get: {