diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 72fc363..a6776a4 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -13,7 +13,21 @@ read-only enforcement, uncertain write outcomes and sensitive-result handling. Test schema changes with referenced request models, path/query arguments and the MCP protocol rather than only calling an internal function. +Favor functionality over compaction. Keep complete small records, expose all +advanced parameters and preserve every sanitized field in detail reads. Test +exact reconstruction from the public index/continuations, including unknown +metadata, Unicode, null/zero/false and changing source data. A shorter first +response is not sufficient evidence of a better agent workflow. Compare complete +tasks and account for schema lookups and follow-up reads; do not claim model task +parity from deterministic protocol tests alone. + +The `read_only_override` catalog field is for reviewed exceptions such as POST +reads. Do not classify an operation as read-only merely to reduce friction: +`check_key_health`, for example, can emit callback test logs. + Release a reviewed commit by updating the package version, running CI, building with `uv build`, and attaching the wheel and source distribution to a matching -GitHub tag. Client integrations should pin a release or commit. PyPI publishing +GitHub tag. The default client launcher tracks `main` and refreshes the connector +on each start, so merged changes do not require a version bump to reach those +clients. Fixed deployments can pin a release or commit instead. PyPI publishing can be added through the organization's trusted publishing configuration. diff --git a/README.md b/README.md index f88ce57..304e23b 100644 --- a/README.md +++ b/README.md @@ -17,8 +17,11 @@ server to a client that supports local MCP processes: "litellm-admin": { "command": "uvx", "args": [ + "--isolated", + "--refresh-package", + "litellm-admin-mcp", "--from", - "git+https://github.com/BerriAI/litellm-admin-mcp.git@v0.1.0", + "git+https://github.com/BerriAI/litellm-admin-mcp.git@main", "litellm-admin-mcp" ], "env": { @@ -33,6 +36,20 @@ server to a client that supports local MCP processes: Use your client's secret storage if available, and keep its configuration private. The connector requires Python 3.12 or later; uv can install a compatible Python. +This configuration automatically checks GitHub's `main` branch whenever your +client starts the MCP server. After a change is merged, restart the MCP connection +or your client to load it; no package version bump or manual reinstall is needed. +`--refresh-package` refreshes this connector's cache, and `--isolated` prevents an +older separately installed tool from taking precedence. See uv's +[tool environments](https://docs.astral.sh/uv/concepts/tools/) and +[cache refresh options](https://docs.astral.sh/uv/concepts/cache/). + +If you already configured a release such as `@v0.1.0`, replace its `args` with the +example above once. Existing pinned configurations cannot update themselves from +a README change. Startup requires access to GitHub; a running process keeps its +current version until restarted. For a fixed deployment, replace `@main` with a +release tag or full commit SHA and omit `--refresh-package` and its package name. + Try: - “List my teams and their budgets.” @@ -62,10 +79,130 @@ ID fails discovery until the connector is updated. Tools have stable names such as `create_key`, `add_model`, `list_teams` and `update_budget`; their arguments are grouped into `body`, `query` and `path` to avoid ambiguous parameter names. -Set `LITELLM_ADMIN_READ_ONLY=true` to expose only GET operations. Optionally set +Set `LITELLM_ADMIN_READ_ONLY=true` to expose only reviewed read operations (including +the POST reads `get_budgets` and `get_organization_legacy`). Optionally set `LITELLM_ADMIN_TOOLS=create_key,list_keys,list_teams` to restrict the catalog. The gateway still enforces each caller's permissions and feature entitlements. +## Efficient results and tool discovery + +The default keeps small results complete and pages large results. All 65 actions, +advanced arguments and previously accessible result fields remain available. +There are two additional read-only tools: `describe_admin_tool` and +`read_admin_result`. + +### Results: complete records, then exact detail reads + +Results up to 16 KB of compact JSON are returned whole. Larger results return an +index with complete small values, useful record previews, unread field counts, +and exact `read` / `next` arguments for `read_admin_result`. This uses structural +paging, not a whitelist of "important" fields or an AI-generated summary. +Unknown metadata, `null`, `0`, `false`, empty arrays and empty objects are preserved. +A preview with `complete: false` is not the entire record; read the relevant +omitted fields before drawing conclusions. For example, a user profile can stay +complete while its embedded keys and teams are read separately. + +To retrieve detail, call `read_admin_result` with the returned `read` arguments. +Follow `next` to enumerate every object field or array item. Paths are JSON +Pointers supplied by the index; string chunks use Unicode code-point offsets. +To retrieve the entire original result in one response, use its `full_read` +arguments (an empty path selects the root): + +```json +{"result_id": "", "path": "", "view": "full"} +``` + +These reads use the same saved result, even if gateway data changes. Reading a +write's result never repeats the write. **Do not repeat a create, rotate or delete +operation just to request a different response view.** Newly created/rotated keys +are always delivered inline in the full result and are never cached. + +Every administrative tool also accepts `response: {"view": "full"}` to request +an inline result on its initial call. Set `LITELLM_ADMIN_RESPONSE_VIEW=full` to make +that the process default. Neither setting changes the gateway query, its scope, +or the existing secret-redaction rules. + +Snapshots are bound to both the caller's credential and current admin identity, +and expire after five minutes. Each read rechecks authorization before releasing +data. The in-memory store admits at most 16 snapshots per credential and 64 MB of +serialized data in total; it never evicts an unexpired result to make room. When +full, or when a result contains recognizable credentials, the original complete +response is returned inline instead. Expired entries are removed on the next +cache access; shutdown clears the store. HTTP deployments with multiple workers +must route subsequent result reads to the same process. A missing/expired result +does not authorize repeating a write. + +Collection pages target 16 KB plus index overhead, up to 20 entries by default; +large individual fields remain explicitly retrievable. Administrative calls +stream at most 16 MB from the gateway in either view, replacing the previous +2 MB limit checked after downloading the entire response. Use upstream pagination +or filters for larger queries. Explicit full reads can be large: request pages +if the client has an output limit. + +### Tool search: use the client's native discovery + +By default, the server publishes ordinary named MCP tools with their complete +input schemas and individual safety annotations. A client such as Codex can use +native tool search/deferred loading with this interface; no custom dispatcher or +second schema lookup is required. The client controls whether definitions enter +model context eagerly or on demand. The size of `tools/list` is not a measurement +of initial model tokens. See OpenAI's [tool search documentation](https://developers.openai.com/api/docs/guides/tools-tool-search). + +For a client that loads every schema upfront, an optional fallback is available: + +```sh +export LITELLM_ADMIN_SCHEMA_MODE=discovery +``` + +This keeps all original tool names, descriptions, required argument groups and +safety annotations, while deferring parameter detail. Before an operation, the +agent calls `describe_admin_tool` with its exact name to obtain **every** supported +argument, constraint, default, example and referenced definition, then calls the +original named tool. Execution still validates the complete gateway schema. +This mode intentionally uses permissive argument declarations; it is not native +tool search and adds a schema lookup. Prefer the default `full` schema mode when +the client already supports native tool search. The helper can also page an +unusually large schema with `response: {"view": "compact"}`. + +The tool allowlist and read-only policy apply to discovery and direct execution +in both modes. Helpers cannot invoke administrative operations or arbitrary URLs. + +### Client compatibility + +Official documentation checked on September 23, 2026. The server uses standard +MCP tools over stdio or Streamable HTTP, so clients control native search and +deferred loading. These are documentation-based compatibility conclusions, not +end-to-end tests inside each application. + +| Client | Documented discovery behavior | Suggested schema mode | +| --- | --- | --- | +| [Claude Code](https://code.claude.com/docs/en/mcp#scale-with-mcp-tool-search) | MCP tool search is on by default for supported models/providers. Tool names and server instructions load first; full definitions load when needed. | Keep the default `full`. | +| [Cursor](https://cursor.com/blog/dynamic-context-discovery#4-efficiently-loading-only-the-mcp-tools-needed) | Dynamic MCP discovery keeps definitions in files and loads relevant tools on demand. | Keep the default `full`. | +| [VS Code / GitHub Copilot](https://code.visualstudio.com/docs/agents/reference/ai-settings) | Experimental `github.copilot.chat.virtualTools.threshold` groups tools for on-demand activation; its documented default threshold is 128 tools. This connector's 67 tools alone do not cross that threshold. | Keep `full` with client-side discovery; consider `discovery` if the client eagerly loads the catalog. | +| [OpenCode](https://opencode.ai/docs/mcp-servers/) | Supports local and remote MCP servers. Its MCP, tools and configuration docs do not establish native tool search and explicitly warn about tool context cost. | `full` works as a standard MCP interface; `discovery` is the fallback to reduce eager schema loading. | + +Claude Code requires a model/provider supporting `tool_reference` blocks. Its +current docs list Claude 4.5-generation models and later. It normally disables +tool search if **model traffic** uses a non-first-party `ANTHROPIC_BASE_URL`. +`ENABLE_TOOL_SEARCH=true` can override that proxy fallback only if the proxy +supports the required blocks; other provider or organization restrictions may +still apply. Pointing this MCP at a LiteLLM gateway with `LITELLM_BASE_URL` is a +separate setting and does not itself change Claude Code's model provider. + +Claude Code recommends concise server instructions explaining the tool categories +and when to search. This server provides those instructions and keeps them below +its documented default 2,048-character truncation limit. Leave `alwaysLoad` unset +for this server if you want Claude Code to defer its tools. + +OpenCode uses an `mcp` configuration map with `type: "local"`, a `command` array +and an `environment` object (rather than the `mcpServers` / `env` example above). +For the lightweight fallback, put `LITELLM_ADMIN_SCHEMA_MODE=discovery` in that +server's `environment`. Remote configurations support bearer `headers` and can +set `oauth: false` for this connector's API-key authentication. + +Compact results and exact detail reads work independently of native tool search; +neither requires client-specific extensions or removes administrative actions. + ## Host an HTTP connector For remote clients, run the same package with Streamable HTTP: @@ -73,10 +210,14 @@ For remote clients, run the same package with Streamable HTTP: ```sh export LITELLM_BASE_URL=https://your-gateway.example.com export LITELLM_MCP_PUBLIC_URL=https://admin-mcp.example.com -uvx --from git+https://github.com/BerriAI/litellm-admin-mcp.git@v0.1.0 \ +uvx --isolated --refresh-package litellm-admin-mcp \ + --from git+https://github.com/BerriAI/litellm-admin-mcp.git@main \ litellm-admin-mcp --transport streamable-http --port 8080 ``` +This command loads the latest `main` on each server start. Restart the hosted +process to apply updates; connected clients can keep the same MCP URL. + Put an HTTPS reverse proxy in front of port 8080. The MCP endpoint is `/mcp` and the process health endpoint is `/healthz`. The public URL configures accepted Host and Origin values. One installation connects to one trusted gateway. @@ -111,6 +252,9 @@ docker run --rm -p 127.0.0.1:8080:8080 \ litellm-admin-mcp ``` +Docker packages the checked-out source into the image. Pull the desired revision, +rebuild the image and redeploy the container to update it. + ## Using it from the LiteLLM Admin Agent The [LiteLLM Admin Agent](https://github.com/BerriAI/litellm-admin-agent) can use @@ -139,6 +283,7 @@ MCP, since endpoint availability depends on the gateway release. uv venv --python 3.12 uv pip install -e '.[test]' uv run --extra test pytest +uv run python scripts/benchmark_efficiency.py uv build docker build -t litellm-admin-mcp:smoke . python scripts/smoke_container.py @@ -147,6 +292,10 @@ python scripts/smoke_container.py Tests use real MCP sessions with instrumented gateway fixtures; they do not call an LLM or modify a production gateway. See [SECURITY.md](SECURITY.md) for the credential boundary and [CONTRIBUTING.md](CONTRIBUTING.md) for adding tools. +The synthetic benchmark verifies exact reconstruction and reports both the first +page and the cost of retrieving every field; optional `tiktoken` installation adds +comparison token counts. It does not measure client startup context or agent +reasoning quality. ## License diff --git a/SECURITY.md b/SECURITY.md index 0787751..6268e81 100644 --- a/SECURITY.md +++ b/SECURITY.md @@ -10,6 +10,22 @@ to start HTTP mode with a process-wide gateway key. The gateway's live proxy-adm role is checked for discovery, before execution and before returning results. The caller's own key remains the execution identity and audit actor. +Large sanitized results can be retained as short-lived in-memory snapshots. +Handles are unguessable and bound to both a keyed digest of the requesting +credential and the verified principal; another admin cannot read them. Each +read checks the current role before and after access. Raw caller credentials +are not stored in the cache. Secret-returning key operations and results with +recognizable credentials bypass it and return their complete response inline. +Snapshot reads never re-execute the original operation. Expiration and capacity +limits are documented in the README; use the same server process for subsequent +reads. The existing result redaction policy is unchanged. + +Full schemas remain the default, including for native client tool search. The +optional lightweight discovery mode retains individual tool names and safety +annotations. It only changes advertised parameter detail: the full live schema +still validates every call. Schema/result helpers are read-only, respect the +operation allowlist and cannot act as a generic execution dispatcher. + Read-only mode and the tool allowlist are enforced by the server, including on direct calls to tools hidden from discovery. Model-returned provider parameters are restricted to model and credential-name identifiers. New virtual keys are diff --git a/scripts/benchmark_efficiency.py b/scripts/benchmark_efficiency.py new file mode 100644 index 0000000..256fb5d --- /dev/null +++ b/scripts/benchmark_efficiency.py @@ -0,0 +1,76 @@ +"""Synthetic response measurements; no gateway, credentials or inference calls. + +Run: python scripts/benchmark_efficiency.py +Optionally install tiktoken to include o200k_base comparison counts. These are +payload measurements, not claims about startup billing or agent task quality. +""" +import json + +from litellm_admin_mcp.results import ResultStore, encode, page + +try: + import tiktoken + tokenizer = tiktoken.get_encoding("o200k_base") +except ImportError: + tokenizer = None + + +keys = [{"id": f"key-{n}", "alias": f"project-{n}", "max_budget": 100, + "spend": n / 10, "blocked": False, "expires": None, + "models": ["example-model"], "metadata": {f"setting-{i}": f"value-{n}-{i}" for i in range(30)}} + for n in range(123)] +fixtures = { + "user_with_relations": {"user_id": "example-admin", "user_info": { + "user_id": "example-admin", "user_role": "proxy_admin", "max_budget": None, + "spend": 0, "blocked": False, "metadata": {"custom_setting": "preserved"}}, + "keys": keys, "teams": [{"id": f"team-{i}", "keys": keys[i:i + 8]} for i in range(14)]}, + "teams_with_keys": [{"id": f"team-{i}", "max_budget": None, "spend": i, + "keys": keys[i:i + 8]} for i in range(14)], + "request_history": {"request_id": "example-request", "request": {"input": [ + {"role": "user", "content": f"Message {i}: " + "Example context. " * 100} for i in range(199)]}, + "response": "Example answer", "metadata": {"rare_debugging_field": False}}, +} + + +def measure(value): + raw = encode(value) + return {"bytes": len(raw), **({"comparison_tokens": len(tokenizer.encode(raw.decode()))} if tokenizer else {})} + + +def reconstruct(item, token, path="", outputs=None): + outputs = outputs if outputs is not None else [] + current = page(item, token, path=path) + typ = current["type"] + result = {} if typ == "object" else [] if typ == "array" else "" if typ == "string" else current["value"] + while True: + outputs.append(current) + if typ in {"object", "array"}: + for entry in current["entries"]: + value = entry["value"] if entry["complete"] else reconstruct(item, token, entry["path"], outputs) + if typ == "object": + result[entry["key"]] = value + else: + result.append(value) + elif typ == "string": + result += current["value"] + if current.get("next") is None: + return result + args = {k: v for k, v in current["next"].items() if k != "result_id"} + current = page(item, token, **args) + + +if __name__ == "__main__": + results = {} + for name, original in fixtures.items(): + store = ResultStore() + token = store.put(name, original, "synthetic-caller", "synthetic-admin") + item = store.get(token, "synthetic-caller", "synthetic-admin") + outputs = [] + assert reconstruct(item, token, outputs=outputs) == original + totals = {key: sum(measure(p)[key] for p in outputs) for key in measure(outputs[0])} + results[name] = {"original": measure(original), "first_page": measure(outputs[0]), + "complete_paged_retrieval": totals, "pages_for_every_field": len(outputs), + "exact_reconstruction": True} + print(json.dumps({"tokenizer": "o200k_base" if tokenizer else None, "synthetic_only": True, + "note": "Reading every field includes index/preview overhead; small first pages do not prove task savings.", + "results": results}, indent=2)) diff --git a/src/litellm_admin_mcp/catalog.py b/src/litellm_admin_mcp/catalog.py index ff3a2d4..6ec758c 100644 --- a/src/litellm_admin_mcp/catalog.py +++ b/src/litellm_admin_mcp/catalog.py @@ -11,10 +11,11 @@ class Operation: path: str operation_id: str summary: str + read_only_override: bool | None = None @property def read_only(self) -> bool: - return self.method == "GET" + return self.method == "GET" if self.read_only_override is None else self.read_only_override OPERATIONS = tuple(Operation(**entry) for entry in json.loads( diff --git a/src/litellm_admin_mcp/config.py b/src/litellm_admin_mcp/config.py index fde0a97..884221e 100644 --- a/src/litellm_admin_mcp/config.py +++ b/src/litellm_admin_mcp/config.py @@ -40,9 +40,15 @@ class Config: read_only: bool = False allowed_tools: frozenset[str] = frozenset() public_url: str = "" + response_view: str = "compact" + schema_mode: str = "full" def __post_init__(self): object.__setattr__(self, "base_url", gateway_origin(self.base_url)) + if self.response_view not in {"full", "compact"}: + raise ValueError("LITELLM_ADMIN_RESPONSE_VIEW must be full or compact") + if self.schema_mode not in {"full", "discovery"}: + raise ValueError("LITELLM_ADMIN_SCHEMA_MODE must be full or discovery") if self.allowed_tools - BY_NAME.keys(): raise ValueError("LITELLM_ADMIN_TOOLS contains unknown tool names; inspect the published tool catalog.") if self.public_url: @@ -58,4 +64,6 @@ def read(cls) -> "Config": read_only=env_bool("LITELLM_ADMIN_READ_ONLY"), allowed_tools=frozenset(x.strip() for x in os.getenv("LITELLM_ADMIN_TOOLS", "").split(",") if x.strip()), public_url=os.getenv("LITELLM_MCP_PUBLIC_URL", "").rstrip("/"), + response_view=os.getenv("LITELLM_ADMIN_RESPONSE_VIEW", "compact").strip(), + schema_mode=os.getenv("LITELLM_ADMIN_SCHEMA_MODE", "full").strip(), ) diff --git a/src/litellm_admin_mcp/gateway.py b/src/litellm_admin_mcp/gateway.py index 3805ad7..a4de0f3 100644 --- a/src/litellm_admin_mcp/gateway.py +++ b/src/litellm_admin_mcp/gateway.py @@ -12,7 +12,20 @@ from .catalog import BY_NAME, OPERATIONS from .config import Config -from .schema import SchemaError, tool_schema +from .schema import SchemaError, discovery_schema, tool_schema +from .results import (MAX_RESULT_BYTES, PAGE_BYTES, READ_SCHEMA, RESPONSE_SCHEMA, + ResultError, ResultStore, encode, page, select) + + +DESCRIBE_SCHEMA = { + "type": "object", "additionalProperties": False, "required": ["name"], + "properties": {"name": {"type": "string", "description": "Exact name of an enabled administrative tool."}, + "response": RESPONSE_SCHEMA}, +} +HELPERS = { + "describe_admin_tool": (DESCRIBE_SCHEMA, "Read the complete input schema and annotations of an enabled tool, including all advanced arguments and referenced definitions. Required before using a tool in discovery schema mode. Returns full detail by default; response.view=compact pages unusually large schemas."), + "read_admin_result": (READ_SCHEMA, "Read an exact saved result without re-executing its operation. Browse fields, array items and string chunks using returned paths/next arguments. view=full returns the entire selected value. Snapshots expire after five minutes; never repeat a write to recover its response."), +} class AdminError(Exception): @@ -66,6 +79,7 @@ def __init__(self, config: Config, client: httpx2.AsyncClient | None = None): self._tools: dict[str, types.Tool] | None = None self._loaded_at = 0.0 self._schema_lock = asyncio.Lock() + self.results = ResultStore() async def __aenter__(self): if self.client is None: @@ -73,26 +87,32 @@ async def __aenter__(self): return self async def __aexit__(self, *args): + self.results.clear() if self._owns_client and self.client is not None: await self.client.aclose() self.client = None - async def _request(self, method: str, path: str, credential: str, **kwargs) -> Any: + async def _request(self, method: str, path: str, credential: str, *, + read_only: bool | None = None, max_bytes: int = 2_000_000, **kwargs) -> Any: assert self.client is not None - write = method != "GET" + write = method != "GET" if read_only is None else not read_only uncertain = " Check gateway state before retrying this change." if write else "" try: - result = await self.client.request(method, self.config.base_url + path, - headers={"Authorization": "Bearer " + valid_credential(credential), "Cookie": "", **kwargs.pop("headers", {})}, - follow_redirects=False, **kwargs) + async with self.client.stream(method, self.config.base_url + path, + headers={"Authorization": "Bearer " + valid_credential(credential), "Cookie": "", **kwargs.pop("headers", {})}, + follow_redirects=False, **kwargs) as result: + if not 200 <= result.status_code < 300: + raise AdminError(f"Gateway returned HTTP {result.status_code}." + uncertain, result.status_code) + chunks, size = [], 0 + async for chunk in result.aiter_bytes(): + size += len(chunk) + if size > max_bytes: + raise AdminError("Result exceeds the upstream byte limit; use pagination or a narrower query." + uncertain) + chunks.append(chunk) except httpx2.RequestError: raise AdminError("Gateway request could not be completed." + uncertain, 504) from None - if not 200 <= result.status_code < 300: - raise AdminError(f"Gateway returned HTTP {result.status_code}." + uncertain, result.status_code) - if len(result.content) > 2_000_000 and path != "/openapi.json": - raise AdminError("Result is too large; use pagination or a narrower query." + uncertain) try: - return result.json() + return json.loads(b"".join(chunks)) except ValueError: raise AdminError("Gateway returned an invalid JSON response." + uncertain, 502) from None @@ -108,10 +128,22 @@ async def authorize(self, credential: str) -> str: async def tools(self, credential: str) -> dict[str, types.Tool]: await self.authorize(credential) + canonical = await self._load_tools(credential) + visible = {} + for name, tool in canonical.items(): + visible[name] = (tool.model_copy(update={"input_schema": discovery_schema(tool.input_schema, name), + "description": tool.description + " Call describe_admin_tool for the complete argument schema before use."}) + if self.config.schema_mode == "discovery" else tool) + for name, (schema, description) in HELPERS.items(): + visible[name] = types.Tool(name=name, description=description, inputSchema=schema, + annotations=types.ToolAnnotations(readOnlyHint=True, destructiveHint=False, idempotentHint=True, openWorldHint=False)) + return visible + + async def _load_tools(self, credential: str) -> dict[str, types.Tool]: async with self._schema_lock: if self._tools is not None and time.monotonic() - self._loaded_at < 300: return self._tools - spec = await self._request("GET", "/openapi.json", credential) + spec = await self._request("GET", "/openapi.json", credential, max_bytes=64_000_000) if not isinstance(spec, dict) or not isinstance(spec.get("paths"), dict): raise AdminError("Gateway OpenAPI schema is unavailable.", 503) selected = {} @@ -128,6 +160,7 @@ async def tools(self, credential: str) -> dict[str, types.Tool]: raise AdminError(f"Gateway operation changed for {item.name}; update the connector before using it.", 503) try: schema = tool_schema(spec, path, operation) + schema["properties"]["response"] = RESPONSE_SCHEMA jsonschema.Draft202012Validator.check_schema(schema) except (SchemaError, jsonschema.SchemaError, TypeError, KeyError): raise AdminError(f"Gateway schema is unsupported for {item.name}.", 503) from None @@ -146,13 +179,65 @@ async def tools(self, credential: str) -> dict[str, types.Tool]: return selected async def call(self, name: str, arguments: dict, credential: str) -> Any: - available = await self.tools(credential) - if name not in available: + # No redundant call to tools(): execution and disclosure each have a + # fresh authorization check. Roles/identities are never cached. + principal = await self.authorize(credential) + loaded_at = self._loaded_at + available = await self._load_tools(credential) + if name not in available and name not in HELPERS: raise AdminError("This tool is not enabled for this connector.", 403) try: - jsonschema.Draft202012Validator(available[name].input_schema).validate(arguments) - except jsonschema.ValidationError: - raise AdminError("Arguments do not match the discovered tool schema.") from None + schema = HELPERS[name][0] if name in HELPERS else available[name].input_schema + jsonschema.Draft202012Validator(schema).validate(arguments) + except jsonschema.ValidationError as exc: + location = "/" + "/".join(str(p) for p in exc.absolute_path) + raise AdminError(f"Arguments do not match the complete tool schema at {location}. Use describe_admin_tool for all accepted fields.") from None + if name in HELPERS: + return await self._helper(name, arguments, credential, principal, available) + if loaded_at != self._loaded_at: + # Discovery awaited upstream I/O; check the original principal + # again before execution, without caching or rebinding identity. + await self._verify(credential, principal) + view = arguments.get("response", {}).get("view", self.config.response_view) + return await self._execute(name, arguments, credential, principal, view) + + async def _verify(self, credential: str, principal: str): + if await self.authorize(credential) != principal: + raise AdminError("Admin identity changed during the request. Inspect gateway state before retrying.", 403) + + def _view(self, name: str, value: Any, credential: str, principal: str, view: str) -> Any: + if view == "compact" and len(encode(value)) > PAGE_BYTES: + token = self.results.put(name, value, credential, principal) + if token is not None: + return page(self.results.get(token, credential, principal), token) + return value + + async def _helper(self, name, arguments, credential, principal, available): + try: + if name == "describe_admin_tool": + target = arguments["name"] + if target not in available: + raise AdminError("This tool is not enabled for this connector.", 403) + result = available[target].model_dump(by_alias=True, exclude_none=True, mode="json") + await self._verify(credential, principal) + return self._view("schema:" + target, result, credential, principal, + arguments.get("response", {}).get("view", "full")) + token = arguments["result_id"] + item = self.results.get(token, credential, principal) + if item.source.removeprefix("schema:") not in available: + raise AdminError("The source operation is no longer enabled.", 403) + if arguments.get("view", "page") == "full": + if any(k in arguments for k in ("offset", "limit", "max_chars")): + raise AdminError("Full reads do not accept pagination arguments.") + result = select(json.loads(item.raw), arguments.get("path", "")) + else: + result = page(item, token, **{k: v for k, v in arguments.items() if k not in {"result_id", "view"}}) + await self._verify(credential, principal) + return result + except ResultError as exc: + raise AdminError(str(exc), 404) from None + + async def _execute(self, name, arguments, credential, principal, view): item = BY_NAME[name] path = item.path for field, value in arguments.get("path", {}).items(): @@ -166,10 +251,16 @@ async def call(self, name: str, arguments: dict, credential: str) -> Any: continue for part in value if isinstance(value, list) else [value]: query.append((field, json.dumps(part) if isinstance(part, (dict, bool)) else str(part))) - principal = await self.authorize(credential) - result = await self._request(item.method, path, credential, params=query, - **({"json": arguments["body"]} if "body" in arguments else {}), - headers={"litellm-changed-by": principal}) - if await self.authorize(credential) != principal: - raise AdminError("Admin identity changed during the request. Inspect gateway state before retrying.", 403) - return safe_result(result, credential) + body, headers = {}, {"litellm-changed-by": principal} + if "body" in arguments: + if arguments["body"] is None: + # httpx's json=None means "no body", not the JSON value null. + body = {"content": b"null"} + headers["Content-Type"] = "application/json" + else: + body = {"json": arguments["body"]} + result = await self._request(item.method, path, credential, params=query, **body, + read_only=item.read_only, max_bytes=MAX_RESULT_BYTES, + headers=headers) + await self._verify(credential, principal) + return self._view(name, safe_result(result, credential), credential, principal, view) diff --git a/src/litellm_admin_mcp/operations.json b/src/litellm_admin_mcp/operations.json index 5964550..bf1bc22 100644 --- a/src/litellm_admin_mcp/operations.json +++ b/src/litellm_admin_mcp/operations.json @@ -24,21 +24,21 @@ "method": "POST", "path": "/key/generate", "operation_id": "generate_key_fn_key_generate_post", - "summary": "Generate Key Fn", + "summary": "Create a virtual key with model access, budgets and rate limits", "name": "create_key" }, { "method": "POST", "path": "/key/service-account/generate", "operation_id": "generate_service_account_key_fn_key_service_account_generate_post", - "summary": "Generate Service Account Key Fn", + "summary": "Create a team-owned service account key", "name": "create_service_account_key" }, { "method": "POST", "path": "/key/update", "operation_id": "update_key_fn_key_update_post", - "summary": "Update Key Fn", + "summary": "Update a virtual key and its access or spending limits", "name": "update_key" }, { @@ -66,21 +66,21 @@ "method": "GET", "path": "/key/info", "operation_id": "info_key_fn_key_info_get", - "summary": "Info Key Fn", + "summary": "Inspect a virtual key, its scope, budget, expiry and configuration", "name": "get_key" }, { "method": "POST", "path": "/key/regenerate", "operation_id": "regenerate_key_fn_key_regenerate_post", - "summary": "Regenerate Key Fn", + "summary": "Rotate a virtual key using a body identifier", "name": "regenerate_key" }, { "method": "POST", "path": "/key/{key}/regenerate", "operation_id": "regenerate_key_fn_key__key__regenerate_post", - "summary": "Regenerate Key Fn", + "summary": "Rotate a virtual key using a path identifier", "name": "regenerate_key_by_id" }, { @@ -94,7 +94,7 @@ "method": "GET", "path": "/key/list", "operation_id": "list_keys_key_list_get", - "summary": "List Keys", + "summary": "List virtual keys with pagination, filters and optional full records", "name": "list_keys" }, { @@ -122,14 +122,14 @@ "method": "POST", "path": "/key/health", "operation_id": "key_health_key_health_post", - "summary": "Key Health", + "summary": "Check the caller credential; this can emit callback test logs", "name": "check_key_health" }, { "method": "GET", "path": "/user/info", "operation_id": "user_info_user_info_get", - "summary": "User Info", + "summary": "Inspect user information; omitting user_id as an admin can include gateway-wide related keys and teams", "name": "get_user" }, { @@ -227,7 +227,7 @@ "method": "GET", "path": "/team/{team_id}/members/me", "operation_id": "team_member_me_team__team_id__members_me_get", - "summary": "Team Member Me", + "summary": "Inspect the caller membership in a selected team", "name": "get_my_team_membership" }, { @@ -255,7 +255,7 @@ "method": "GET", "path": "/team/available", "operation_id": "list_available_teams_team_available_get", - "summary": "List Available Teams", + "summary": "List teams available in the caller membership scope", "name": "list_available_teams" }, { @@ -346,8 +346,9 @@ "method": "POST", "path": "/organization/info", "operation_id": "deprecated_info_organization_organization_info_post", - "summary": "Deprecated Info Organization", - "name": "get_organization_legacy" + "summary": "Read organization details through the legacy POST endpoint", + "name": "get_organization_legacy", + "read_only_override": true }, { "method": "GET", @@ -430,8 +431,9 @@ "method": "POST", "path": "/budget/info", "operation_id": "info_budget_budget_info_post", - "summary": "Info Budget", - "name": "get_budgets" + "summary": "Read selected budget records", + "name": "get_budgets", + "read_only_override": true }, { "method": "GET", diff --git a/src/litellm_admin_mcp/results.py b/src/litellm_admin_mcp/results.py new file mode 100644 index 0000000..92ddc3b --- /dev/null +++ b/src/litellm_admin_mcp/results.py @@ -0,0 +1,242 @@ +"""Bounded, caller-bound snapshots of already sanitized tool results. + +Views are structural, not field whitelists or model-generated summaries. Every +omitted value has an exact read path; small objects are returned whole. +""" +from collections import OrderedDict +from dataclasses import dataclass +import hashlib +import hmac +import json +import re +import secrets +import time +from typing import Any + + +PAGE_BYTES = 16_000 +ITEM_BYTES = PAGE_BYTES - 1000 +PREVIEW_BYTES = 6_000 +MAX_RESULT_BYTES = 16_000_000 +SECRET_RESULTS = {"create_key", "create_service_account_key", "regenerate_key", "regenerate_key_by_id"} + + +class ResultError(ValueError): + pass + + +def encode(value: Any) -> bytes: + return json.dumps(value, ensure_ascii=False, separators=(",", ":")).encode("utf-8") + + +def pointer(parent: str, key: str | int) -> str: + return parent + "/" + str(key).replace("~", "~0").replace("/", "~1") + + +def select(value: Any, path: str) -> Any: + if not path: + return value + if not path.startswith("/") or re.search(r"~(?![01])", path): + raise ResultError("Use a JSON Pointer path returned by the result index.") + try: + for part in path[1:].split("/"): + part = part.replace("~1", "/").replace("~0", "~") + if isinstance(value, list): + if not re.fullmatch(r"0|[1-9][0-9]*", part): + raise ValueError + value = value[int(part)] + elif isinstance(value, dict): + value = value[part] + else: + raise ValueError + except (KeyError, IndexError, ValueError): + raise ResultError("The result does not contain that path; browse its parent first.") from None + return value + + +def kind(value: Any) -> str: + if value is None: + return "null" + if isinstance(value, bool): + return "boolean" + if isinstance(value, dict): + return "object" + if isinstance(value, list): + return "array" + if isinstance(value, str): + return "string" + return "number" + + +def contains_credential(value: Any, *, source: str = "", path: tuple[str | int, ...] = ()) -> bool: + if isinstance(value, dict): + for key, child in value.items(): + # Model-info reads use data[i].model_info.key for a catalog identifier. + # Exempt only that field's name; still scan its value for credentials. + catalog_key = (source in {"list_models", "get_model"} and len(path) == 3 + and path[0] == "data" and isinstance(path[1], int) + and path[2] == "model_info" and key == "key") + if (not catalog_key + and key.lower() in {"key", "api_key", "token", "access_token", "refresh_token", "password", + "secret", "client_secret", "authorization", "authentication_token"} + and isinstance(child, str) and child + and child not in {"[redacted]", "[credential redacted]"} + and not re.fullmatch(r"[a-f0-9]{64}", child)): + return True + if contains_credential(child, source=source, path=(*path, key)): + return True + elif isinstance(value, list): + return any(contains_credential(v, source=source, path=(*path, i)) for i, v in enumerate(value)) + elif isinstance(value, str): + if re.search(r"\bsk-[A-Za-z0-9_+-]+", value): + return True + try: + nested = json.loads(value) + except ValueError: + return False + if isinstance(nested, (dict, list)): + return contains_credential(nested, source=source, path=path) + return False + + +@dataclass(frozen=True) +class Snapshot: + owner: bytes + principal: str + source: str + raw: bytes + deadline: float + expires_at: float + + +class ResultStore: + def __init__(self, *, ttl: int = 300, max_bytes: int = 64_000_000, max_per_owner: int = 16): + self.ttl, self.max_bytes, self.max_per_owner = ttl, max_bytes, max_per_owner + self._salt = secrets.token_bytes(32) + self._items: OrderedDict[str, Snapshot] = OrderedDict() + self._bytes = 0 + + def _owner(self, credential: str) -> bytes: + return hmac.digest(self._salt, credential.encode(), hashlib.sha256) + + def _expire(self): + now = time.monotonic() + for token, item in list(self._items.items()): + if item.deadline <= now: + self._bytes -= len(item.raw) + del self._items[token] + + def clear(self): + self._items.clear() + self._bytes = 0 + + def put(self, source: str, value: Any, credential: str, principal: str) -> str | None: + self._expire() + raw = encode(value) + # New credentials are delivered in the original full response, never + # retained for later reads. Also exclude recognizable keys in log data. + if source in SECRET_RESULTS or contains_credential(value, source=source): + return None + owner = self._owner(credential) + if (len(raw) > MAX_RESULT_BYTES or self._bytes + len(raw) > self.max_bytes + or sum(s.owner == owner for s in self._items.values()) >= self.max_per_owner): + # Do not evict an unexpired snapshot or lose a committed write's + # response. The caller falls back to the complete inline result. + return None + token = secrets.token_urlsafe(24) + self._items[token] = Snapshot(owner, principal, source, raw, + time.monotonic() + self.ttl, time.time() + self.ttl) + self._bytes += len(raw) + return token + + def get(self, token: str, credential: str, principal: str) -> Snapshot: + self._expire() + item = self._items.get(token) + if item is None or item.owner != self._owner(credential) or item.principal != principal: + raise ResultError("Result unavailable to this caller or expired. Do not repeat a write to retrieve its response.") + return item + + +def page(item: Snapshot, token: str, *, path: str = "", offset: int = 0, + limit: int = 20, max_chars: int = 4000) -> dict: + value = select(json.loads(item.raw), path) + node_type = kind(value) + result = {"format": "litellm-admin-page-v1", "result_id": token, "source_tool": item.source, + "expires_at": item.expires_at, "path": path, "type": node_type, + "read_tool": "read_admin_result", "full_read": {"result_id": token, "path": path, "view": "full"}} + if not isinstance(value, (dict, list, str)): + if offset: + raise ResultError("Scalar values do not have offsets.") + return {**result, "value": value, "complete": True} + total = len(value) + if offset > total: + raise ResultError("Offset exceeds this node's length.") + result.update(total=total, offset=offset) + if isinstance(value, str): + # Offsets count Unicode code points, not UTF-8 bytes or tokens. + chunk = value[offset:offset + max_chars] + result.update(value=chunk, offset_unit="unicode_code_points") + end = offset + len(chunk) + result["complete"] = offset == 0 and end == total + else: + entries, used = [], 0 + keys = list(value) if isinstance(value, dict) else range(total) + for key in keys[offset:offset + limit]: + child = value[key] + child_path = pointer(path, key) + entry = {"key" if isinstance(value, dict) else "index": key, + "path": child_path, "type": kind(child)} + raw_size = len(encode(child)) + if raw_size <= ITEM_BYTES: + entry.update(value=child, complete=True) + else: + if entries and PAGE_BYTES - used < PREVIEW_BYTES + 1000: + break + entry.update(complete=False, read={"result_id": token, "path": child_path}) + if isinstance(child, (dict, list, str)): + entry["total"] = len(child) + if isinstance(child, dict): + # Preserve whole small fields, including unknown metadata, + # null/0/false. The preview never pretends to be complete. + preview, omitted, remaining = {}, [], PREVIEW_BYTES + for field, val in child.items(): + cost = len(encode({field: val})) + if cost <= remaining: + preview[field] = val + remaining -= cost + else: + omitted.append(field) + entry.update(preview=preview, omitted_fields=omitted[:20], omitted_count=len(omitted)) + cost = len(encode(entry)) + if entries and used + cost > PAGE_BYTES: + break + entries.append(entry) + used += cost + result["entries"] = entries + end = offset + len(entries) + result["complete"] = offset == 0 and end == total and all(e["complete"] for e in entries) + result["has_more"] = end < total + result["next"] = ({"result_id": token, "path": path, "offset": end, + "limit": limit, "max_chars": max_chars} if end < total else None) + return result + + +RESPONSE_SCHEMA = { + "type": "object", "additionalProperties": False, + "properties": {"view": {"type": "string", "enum": ["full", "compact"], + "description": "full preserves the original result. compact keeps small results whole and pages large results without discarding data. Never repeat a write for a different view; use read_admin_result."}}, +} + +READ_SCHEMA = { + "type": "object", "additionalProperties": False, "required": ["result_id"], + "properties": { + "result_id": {"type": "string"}, + "path": {"type": "string", "default": "", "description": "Exact JSON Pointer from an index; empty string selects the entire original result."}, + "view": {"type": "string", "enum": ["page", "full"], "default": "page"}, + "offset": {"type": "integer", "minimum": 0, "default": 0}, + "limit": {"type": "integer", "minimum": 1, "maximum": 100, "default": 20, + "description": "Object fields or array items per page."}, + "max_chars": {"type": "integer", "minimum": 1, "maximum": 16000, "default": 4000, + "description": "Unicode code points per string chunk; offset uses the same units."}, + }, +} diff --git a/src/litellm_admin_mcp/schema.py b/src/litellm_admin_mcp/schema.py index 5b77c6e..5dd7d85 100644 --- a/src/litellm_admin_mcp/schema.py +++ b/src/litellm_admin_mcp/schema.py @@ -82,6 +82,48 @@ def rewrite(value): return result result = rewrite(root) + # Some clients render a referenced top-level body as `unknown`. Inline that + # one level, preserving every field and constraint, then retain only the + # definitions still needed (including recursive references). + body = result["properties"].get("body", {}) + if set(body) == {"$ref"}: + name = body["$ref"].removeprefix("#/$defs/").replace("~1", "/").replace("~0", "~") + result["properties"]["body"] = deepcopy(definitions[name]) + reachable = set() + + def visit(value): + if isinstance(value, dict): + ref = value.get("$ref") + if isinstance(ref, str) and ref.startswith("#/$defs/"): + name = ref.removeprefix("#/$defs/").replace("~1", "/").replace("~0", "~") + if name not in reachable: + reachable.add(name) + visit(definitions[name]) + for child in value.values(): + visit(child) + elif isinstance(value, list): + for child in value: + visit(child) + + visit(result) + definitions = {name: value for name, value in definitions.items() if name in reachable} if definitions: result["$defs"] = definitions return result + + +def discovery_schema(schema: dict, name: str) -> dict: + """Keep named tools/approval boundaries while deferring parameter detail. + + This fallback declaration accepts the same body/query/path shapes, including + scalar or array bodies. Execution ALWAYS validates the complete schema. + No constraints are guessed from a small subset of commonly used fields. + """ + result = {"type": "object", "additionalProperties": False, "properties": {}} + for group in schema["properties"]: + result["properties"][group] = (deepcopy(schema["properties"][group]) if group == "response" else { + "description": f"Before calling {name}, use describe_admin_tool with name={name} for the complete {group} schema. All original fields remain supported and are validated on execution." + }) + if schema.get("required"): + result["required"] = list(schema["required"]) + return result diff --git a/src/litellm_admin_mcp/server.py b/src/litellm_admin_mcp/server.py index b7d5ce3..868cb4f 100644 --- a/src/litellm_admin_mcp/server.py +++ b/src/litellm_admin_mcp/server.py @@ -39,13 +39,19 @@ async def list_tools(context, params): async def call_tool(context, params): try: result = await gateway.call(params.name, params.arguments or {}, credential(context)) - return types.CallToolResult(content=[types.TextContent(type="text", text=json.dumps(result))]) + return types.CallToolResult(content=[types.TextContent(type="text", text=json.dumps(result, ensure_ascii=False, separators=(",", ":")))]) except AdminError as exc: return types.CallToolResult(isError=True, content=[types.TextContent(type="text", text=str(exc))]) return Server("litellm-admin-mcp", version=__version__, - instructions="Administer the connected LiteLLM gateway only when the user requests it. " + instructions="Use these tools for the connected LiteLLM gateway's virtual keys, teams and members, budgets, " + "model deployments, spending, activity and request logs. Search for these tools when the user asks to inspect " + "or administer those resources. A current proxy-admin identity is required. Make changes only when requested. " "Discover exact tool schemas; arguments are grouped under body, query and path. " + "In discovery schema mode, call describe_admin_tool before using an operation; it returns ALL supported arguments. " + "response.view=compact keeps small results whole and indexes large results. An incomplete preview is not an empty or complete result. " + "Use read_admin_result with returned paths/next arguments for every relevant omitted field; full reads preserve all original data. " + "Never re-execute a write to change its response view. New virtual keys are always delivered inline. " "Look up resource identifiers before making changes and verify changes afterward. " "Use named gateway credentials or environment references for provider authentication. " "Do not request provider secrets in chat. A newly created virtual key is sensitive: deliver it only to its requester. " diff --git a/tests/test_connector.py b/tests/test_connector.py index 1ae2ccf..8faf031 100644 --- a/tests/test_connector.py +++ b/tests/test_connector.py @@ -8,7 +8,7 @@ from conftest import SPEC, data, session_for from litellm_admin_mcp.catalog import OPERATIONS from litellm_admin_mcp.config import Config -from litellm_admin_mcp.gateway import Gateway +from litellm_admin_mcp.gateway import Gateway, HELPERS from litellm_admin_mcp.schema import tool_schema from litellm_admin_mcp.server import create_http_app @@ -16,7 +16,7 @@ async def test_real_mcp_create_key_and_read_back(stub): async with session_for(stub) as session: tools = (await session.list_tools()).tools - assert {t.name for t in tools} == {"create_key", "list_keys", "add_model", "get_model", "reset_key_spend"} + assert {t.name for t in tools} == {"create_key", "list_keys", "add_model", "get_model", "reset_key_spend"} | HELPERS.keys() created = data(await session.call_tool("create_key", {"body": {"key_alias": "engineering", "max_budget": 100}})) assert created["key"].startswith("sk-created-for-") assert created["echo"] == "[credential redacted]" diff --git a/tests/test_progressive.py b/tests/test_progressive.py new file mode 100644 index 0000000..aeb0d60 --- /dev/null +++ b/tests/test_progressive.py @@ -0,0 +1,359 @@ +import copy +import json +import re + +import httpx2 +import jsonschema +import pytest + +from conftest import SPEC, connected, data, session_for +from litellm_admin_mcp.catalog import OPERATIONS +from litellm_admin_mcp.config import Config +from litellm_admin_mcp.gateway import AdminError, Gateway, HELPERS +from litellm_admin_mcp.results import ResultStore + + +async def test_discovery_keeps_named_tools_annotations_and_complete_validation(stub): + async with session_for(stub, schema_mode="discovery") as session: + listed = {t.name: t for t in (await session.list_tools()).tools} + assert not listed["create_key"].annotations.read_only_hint + assert listed["create_key"].annotations.destructive_hint + assert listed["list_keys"].annotations.read_only_hint + detail = data(await session.call_tool("describe_admin_tool", {"name": "create_key"})) + jsonschema.Draft202012Validator(detail["inputSchema"]).validate({"body": {"key_alias": "a", "max_budget": 0}}) + rejected = await session.call_tool("create_key", {"body": {"key_alias": "a", "unsupported": True}}) + assert rejected.is_error and "complete tool schema" in rejected.content[0].text + assert not stub.writes + result = data(await session.call_tool("create_key", {"body": {"key_alias": "a", "max_budget": 0}})) + assert result["max_budget"] == 0 + + +async def test_full_schema_exposes_body_fields_without_losing_nested_definitions(stub): + async with session_for(stub) as session: + tools = {t.name: t for t in (await session.list_tools()).tools} + key_schema = tools["create_key"].input_schema + assert key_schema["properties"]["body"]["type"] == "object" + assert "key_alias" in key_schema["properties"]["body"]["properties"] + assert "$defs" not in key_schema + model_schema = tools["add_model"].input_schema + assert "Params" in model_schema["$defs"] + assert "Deployment" not in model_schema["$defs"] + + +async def test_schema_itself_can_be_paged_and_retrieved_in_full(stub): + stub.spec = copy.deepcopy(SPEC) + stub.spec["components"]["schemas"]["Key"]["properties"]["key_alias"]["description"] = "Long but useful documentation. " * 2000 + async with session_for(stub, schema_mode="discovery") as session: + full = data(await session.call_tool("describe_admin_tool", {"name": "create_key"})) + indexed = data(await session.call_tool("describe_admin_tool", {"name": "create_key", "response": {"view": "compact"}})) + assert data(await session.call_tool("read_admin_result", indexed["full_read"])) == full + + +@pytest.mark.parametrize("schema,value", [ + ({"type": "integer"}, 0), ({"type": "boolean"}, False), ({"type": "null"}, None), + ({"type": "array", "items": {"type": "string"}}, ["uncommon", "argument"]), +]) +async def test_discovery_preserves_non_object_body_shapes(stub, schema, value): + stub.spec = copy.deepcopy(SPEC) + stub.spec["components"]["schemas"]["Key"] = schema + original = stub.handle + seen = [] + def handle(request): + if request.url.path == "/key/generate": + seen.append(json.loads(request.content)) + return httpx2.Response(200, json={"accepted": seen[-1]}) + return original(request) + stub.handle = handle + async with session_for(stub, schema_mode="discovery") as session: + detail = data(await session.call_tool("describe_admin_tool", {"name": "create_key"})) + jsonschema.Draft202012Validator(detail["inputSchema"]).validate({"body": value}) + assert data(await session.call_tool("create_key", {"body": value})) == {"accepted": value} + assert seen == [value] + + +async def test_post_reads_are_available_in_read_only_mode_but_health_is_not(stub): + stub.spec = copy.deepcopy(SPEC) + for op in OPERATIONS: + if op.name in {"get_budgets", "get_organization_legacy", "check_key_health"}: + stub.spec["paths"][op.path] = {op.method.lower(): {"operationId": op.operation_id}} + original = stub.handle + def handle(request): + if request.url.path in {"/budget/info", "/organization/info"}: + return httpx2.Response(200, json=[]) + return original(request) + stub.handle = handle + async with session_for(stub, read_only=True, schema_mode="discovery") as session: + tools = {t.name: t for t in (await session.list_tools()).tools} + assert "check_key_health" not in tools + for name in ("get_budgets", "get_organization_legacy"): + assert tools[name].annotations.read_only_hint + assert not tools[name].annotations.destructive_hint + assert data(await session.call_tool(name, {})) == [] + + +async def test_full_mode_is_unchanged_and_compact_pages_do_not_refetch(stub): + stub.keys = [{"id": i, "max_budget": None, "spend": 0, "blocked": False, + "extra": {"rare_setting": "z" * 8000}} for i in range(15)] + async with session_for(stub, response_view="full") as session: + original = data(await session.call_tool("list_keys", {})) + assert original == {"keys": stub.keys} + paged = data(await session.call_tool("list_keys", {"response": {"view": "compact"}})) + assert paged["format"] == "litellm-admin-page-v1" and not paged["complete"] + count = sum(r.url.path == "/key/list" for r in stub.requests) + stub.keys.clear() + restored = data(await session.call_tool("read_admin_result", paged["full_read"])) + assert restored == original + assert sum(r.url.path == "/key/list" for r in stub.requests) == count + + +async def test_compact_default_keeps_small_results_and_full_override(stub): + async with session_for(stub, response_view="compact") as session: + assert data(await session.call_tool("list_keys", {})) == {"keys": []} + stub.keys = [{"id": i, "metadata": "a" * 1000} for i in range(40)] + assert "result_id" in data(await session.call_tool("list_keys", {})) + assert data(await session.call_tool("list_keys", {"response": {"view": "full"}})) == {"keys": stub.keys} + + +@pytest.mark.parametrize("name", ["list_models", "get_model"]) +async def test_model_catalog_identifiers_allow_lossless_paging(stub, name): + operation = next(op for op in OPERATIONS if op.name == name) + stub.spec = copy.deepcopy(SPEC) + stub.spec["paths"][operation.path] = {"get": {"operationId": operation.operation_id}} + models = [{"model_name": f"deployment-{i}", "model_info": { + "id": f"model-{i}", "key": f"openai/example-model-{i}", + "input_cost_per_token": 0, "supports_vision": False, + "metadata": {"uncommon": None, "details": "x" * 2000}, + }} for i in range(40)] + expected = {"data": models, "total_count": len(models), "current_page": 1, "total_pages": 1} + original = stub.handle + + def handle(request): + if request.url.path == operation.path: + stub.requests.append(request) + return httpx2.Response(200, json=expected) + return original(request) + + stub.handle = handle + async with session_for(stub) as session: + first = data(await session.call_tool(name, {})) + assert first["format"] == "litellm-admin-page-v1" + assert len(json.dumps(first)) < 20000 + assert data(await session.call_tool("read_admin_result", first["full_read"])) == expected + records = next(e for e in first["entries"] if e["key"] == "data") + current = data(await session.call_tool("read_admin_result", records["read"])) + restored = [] + while True: + for entry in current["entries"]: + assert entry["complete"] + restored.append(entry["value"]) + if not current["next"]: + break + current = data(await session.call_tool("read_admin_result", current["next"])) + assert restored == models + assert sum(r.url.path == operation.path for r in stub.requests) == 1 + assert data(await session.call_tool(name, {"response": {"view": "full"}})) == expected + + +async def test_full_write_result_can_be_read_without_repeating_write(stub): + original_handler = stub.handle + payload = {"status": "ok", "previous": {"spend": 20, "custom": "x" * 22000}} + def handle(request): + response = original_handler(request) + return httpx2.Response(200, json=payload) if request.url.path.endswith("/reset_spend") else response + stub.handle = handle + async with session_for(stub, response_view="compact") as session: + result = data(await session.call_tool("reset_key_spend", {"path": {"key": "a" * 64}})) + full = data(await session.call_tool("read_admin_result", result["full_read"])) + assert full == payload + assert len(stub.writes) == 1 + + +async def test_large_new_key_response_is_delivered_inline_and_never_cached(stub): + original = stub.handle + def handle(request): + response = original(request) + if request.url.path == "/key/generate": + return httpx2.Response(200, json={**response.json(), "metadata": "x" * 20000}) + return response + async with httpx2.AsyncClient(transport=httpx2.MockTransport(handle)) as client: + gateway = Gateway(Config("https://gateway.example.com", api_key="alice-admin", response_view="compact"), client) + async with connected(gateway) as session: + result = data(await session.call_tool("create_key", {"body": {"key_alias": "new"}})) + assert result["key"].startswith("sk-created") + assert len(result["metadata"]) == 20000 + assert not gateway.results._items + assert len(stub.writes) == 1 + + +async def test_result_reads_require_same_credential_principal_and_current_role(stub): + stub.keys = [{"data": "x" * 20000}] + async with httpx2.AsyncClient(transport=httpx2.MockTransport(stub.handle)) as client: + gateway = Gateway(Config("https://gateway.example.com", response_view="compact"), client) + result = await gateway.call("list_keys", {}, "alice-admin") + for credential in ("bob-admin", "viewer"): + with pytest.raises(AdminError): + await gateway.call("read_admin_result", result["full_read"], credential) + stub.revoked.add("alice-admin") + with pytest.raises(AdminError): + await gateway.call("read_admin_result", result["full_read"], "alice-admin") + + +async def test_revocation_during_result_read_suppresses_snapshot(stub): + stub.keys = [{"data": "private" * 4000}] + async with httpx2.AsyncClient(transport=httpx2.MockTransport(stub.handle)) as client: + gateway = Gateway(Config("https://gateway.example.com", response_view="compact"), client) + result = await gateway.call("list_keys", {}, "alice-admin") + original_get = gateway.results.get + def get(*args): + item = original_get(*args) + stub.revoked.add("alice-admin") + return item + gateway.results.get = get + with pytest.raises(AdminError, match="proxy admins"): + await gateway.call("read_admin_result", result["full_read"], "alice-admin") + + +async def test_schema_helper_respects_allowlist_and_cannot_execute_hidden_operations(stub): + async with session_for(stub, allowed_tools=frozenset({"list_keys"}), schema_mode="discovery") as session: + for name in ("create_key", "flush_everything", "read_admin_result"): + assert (await session.call_tool("describe_admin_tool", {"name": name})).is_error + assert (await session.call_tool("read_admin_result", {"result_id": "guess", "method": "POST"})).is_error + assert (await session.call_tool("create_key", {"body": {"key_alias": "a"}})).is_error + assert not stub.writes + + +async def test_warm_call_retains_pre_post_checks_without_discovery_auth(stub): + async with session_for(stub) as session: + await session.list_tools() + stub.requests.clear() + await session.call_tool("list_keys", {}) + assert [r.url.path for r in stub.requests] == ["/user/info", "/key/list", "/user/info"] + + +async def test_identity_change_during_cold_discovery_prevents_write(stub): + original = stub.handle + changed = False + def handle(request): + nonlocal changed + result = original(request) + if request.url.path == "/openapi.json": + changed = True + elif request.url.path == "/user/info" and changed: + return httpx2.Response(200, json={"user_id": "new-identity", "user_info": {"user_id": "new-identity", "user_role": "proxy_admin"}}) + return result + stub.handle = handle + async with session_for(stub) as session: + result = await session.call_tool("create_key", {"body": {"key_alias": "a"}}) + assert result.is_error and "identity changed" in result.content[0].text + assert not stub.writes + + +async def test_capacity_falls_back_to_full_write_response_without_retry(stub): + original = stub.handle + payload = {"status": "ok", "previous": "x" * 20000} + def handle(request): + result = original(request) + return httpx2.Response(200, json=payload) if request.url.path.endswith("/reset_spend") else result + async with httpx2.AsyncClient(transport=httpx2.MockTransport(handle)) as client: + gateway = Gateway(Config("https://gateway.example.com", response_view="compact"), client) + gateway.results = ResultStore(max_bytes=1) + assert await gateway.call("reset_key_spend", {"path": {"key": "fixture"}}, "alice-admin") == payload + assert len(stub.writes) == 1 + + +async def test_views_support_large_results_and_stop_at_stream_limit(stub): + original = stub.handle + class Chunks(httpx2.AsyncByteStream): + def __init__(self): + self.count = 0 + async def __aiter__(self): + for _ in range(30): + self.count += 1 + yield b"x" * 1_000_000 + stream = Chunks() + def handle(request): + if request.url.path == "/key/list": + if request.url.params.get("page") == "2": + return httpx2.Response(200, stream=stream) + return httpx2.Response(200, json={"data": "x" * 2_100_000}) + return original(request) + stub.handle = handle + async with session_for(stub) as session: + assert len(data(await session.call_tool("list_keys", {"response": {"view": "full"}}))["data"]) == 2_100_000 + result = data(await session.call_tool("list_keys", {"response": {"view": "compact"}})) + assert result["entries"][0]["total"] == 2_100_000 + error = await session.call_tool("list_keys", {"query": {"page": 2}, "response": {"view": "compact"}}) + assert error.is_error and "byte limit" in error.content[0].text + assert stream.count == 17 + + +async def test_defaults_support_native_tool_search_and_compact_large_results(stub, monkeypatch): + monkeypatch.setenv("LITELLM_BASE_URL", "https://gateway.example.com") + monkeypatch.delenv("LITELLM_ADMIN_RESPONSE_VIEW", raising=False) + monkeypatch.delenv("LITELLM_ADMIN_SCHEMA_MODE", raising=False) + config = Config.read() + assert config.response_view == "compact" and config.schema_mode == "full" + stub.keys = [{"data": "x" * 20000}] + async with session_for(stub) as session: + tools = {t.name: t for t in (await session.list_tools()).tools} + assert "key_alias" in tools["create_key"].input_schema["properties"]["body"]["properties"] + assert "result_id" in data(await session.call_tool("list_keys", {})) + + +@pytest.mark.parametrize("mode", ["full", "discovery"]) +async def test_every_reviewed_operation_retains_mapping_advanced_inputs_and_result(mode): + spec = {"paths": {}, "components": {"schemas": { + "Advanced": {"type": "object", "properties": { + "model_limits": {"type": "object", "additionalProperties": {"type": ["number", "null"]}}, + "enabled": {"type": "boolean"}, "tags": {"type": "array", "items": {"type": "string"}}, + }, "required": ["enabled"], "additionalProperties": False}, + "Body": {"type": "object", "properties": {"advanced": {"$ref": "#/components/schemas/Advanced"}}, + "required": ["advanced"], "additionalProperties": False}, + }}} + for op in OPERATIONS: + parameters = [{"in": "path", "name": name, "required": True, "schema": {"type": "string"}} + for name in re.findall(r"\{([^}]+)\}", op.path)] + parameters += [{"in": "query", "name": "filter", "schema": {"type": "array", "items": {"type": "string"}}}] + definition = {"operationId": op.operation_id, "parameters": parameters} + if op.method != "GET": + definition["requestBody"] = {"required": True, "content": {"application/json": {"schema": {"$ref": "#/components/schemas/Body"}}}} + spec["paths"].setdefault(op.path, {})[op.method.lower()] = definition + calls = [] + result = {"budget": None, "spend": 0, "blocked": False, "metadata": {"unexpected_field": [1, "a", None]}} + def handle(request): + if request.url.path == "/openapi.json": + return httpx2.Response(200, json=spec) + if request.url.path == "/user/info" and "litellm-changed-by" not in request.headers: + return httpx2.Response(200, json={"user_id": "admin", "user_info": {"user_id": "admin", "user_role": "proxy_admin"}}) + calls.append(request) + return httpx2.Response(200, json=result) + async with httpx2.AsyncClient(transport=httpx2.MockTransport(handle)) as client: + gateway = Gateway(Config("https://gateway.example.com", api_key="admin", schema_mode=mode), client) + async with connected(gateway) as session: + listed = {t.name: t for t in (await session.list_tools()).tools} + assert set(listed) - HELPERS.keys() == {op.name for op in OPERATIONS} + for op in OPERATIONS: + detail = data(await session.call_tool("describe_admin_tool", {"name": op.name})) + assert detail["annotations"]["readOnlyHint"] == op.read_only + arguments = {"query": {"filter": ["first", "second"]}} + path_args = {name: "fixture" for name in re.findall(r"\{([^}]+)\}", op.path)} + if path_args: + arguments["path"] = path_args + if op.method != "GET": + arguments["body"] = {"advanced": {"enabled": False, "model_limits": {"rare/model": None, "free/model": 0}, "tags": []}} + jsonschema.Draft202012Validator(detail["inputSchema"]).validate(arguments) + assert data(await session.call_tool(op.name, arguments)) == result + request = calls[-1] + assert request.method == op.method + assert request.url.path == re.sub(r"\{[^}]+\}", "fixture", op.path) + assert request.url.params.get_list("filter") == ["first", "second"] + assert request.headers["litellm-changed-by"] == "admin" + if "body" in arguments: + assert json.loads(request.content) == arguments["body"] + assert len(calls) == 65 + + +@pytest.mark.parametrize("config", [{"schema_mode": "unknown"}, {"response_view": "unknown"}]) +def test_unknown_modes_fail_instead_of_silently_changing_behavior(config): + with pytest.raises(ValueError): + Config("https://gateway.example.com", **config) diff --git a/tests/test_results.py b/tests/test_results.py new file mode 100644 index 0000000..9560886 --- /dev/null +++ b/tests/test_results.py @@ -0,0 +1,142 @@ +import json + +import pytest + +from litellm_admin_mcp.results import ResultError, ResultStore, encode, page, select + + +def rebuild(item, token, path=""): + """Use only the public index/continuations, without knowing field names.""" + first = page(item, token, path=path, limit=3, max_chars=257) + result = {} if first["type"] == "object" else [] if first["type"] == "array" else "" if first["type"] == "string" else first["value"] + current = first + while True: + if first["type"] in {"object", "array"}: + for entry in current["entries"]: + value = entry["value"] if entry["complete"] else rebuild(item, token, entry["read"]["path"]) + if isinstance(result, dict): + result[entry["key"]] = value + else: + result.append(value) + elif first["type"] == "string": + result += current["value"] + if not current.get("next"): + return result + args = dict(current["next"]) + args.pop("result_id") + current = page(item, token, **args) + + +def test_index_and_pages_reconstruct_every_value_without_prior_field_knowledge(): + original = { + "user_info": {"role": "proxy_admin", "max_budget": None, "spend": 0, "blocked": False}, + "keys": [{"id": n, "rare_setting": {"model/a~b": n}, "metadata": {"new_field": "x" * 17100}} + for n in range(17)], + "a/b~c": {"": "🎉e\u0301" * 2700, "empty_object": {}, "empty_array": [], "empty_string": ""}, + "0": [True, False, None, 0, -1, 1.25], + } + store = ResultStore() + token = store.put("get_user", original, "alice", "user-1") + item = store.get(token, "alice", "user-1") + first = page(item, token) + profile = next(e for e in first["entries"] if e["key"] == "user_info") + assert profile["complete"] and profile["value"] == original["user_info"] + assert rebuild(item, token) == original + assert json.loads(item.raw) == original + assert len(encode(first)) < len(item.raw) / 4 + + +def test_record_previews_include_all_small_fields_including_uncommon_ones(): + value = [{"id": "team-1", "custom_policy": False, "budget": None, "spend": 0, + "new_config": {"special": True}, "keys": ["x" * 500 for _ in range(100)]}] + store = ResultStore() + token = store.put("list_teams", value, "alice", "user-1") + result = page(store.get(token, "alice", "user-1"), token) + entry = result["entries"][0] + assert not entry["complete"] + assert entry["preview"] == {k: v for k, v in value[0].items() if k != "keys"} + assert entry["omitted_fields"] == ["keys"] + + +def test_page_boundary_moves_whole_records_to_next_page_instead_of_fragmenting(): + value = [{"id": n, "data": "x" * 7000} for n in range(7)] + store = ResultStore() + token = store.put("list_keys", value, "alice", "user-1") + item = store.get(token, "alice", "user-1") + current, seen, calls = page(item, token), [], 0 + while True: + calls += 1 + for entry in current["entries"]: + assert entry["complete"] + seen.append(entry["value"]) + if not current["next"]: + break + args = {k: v for k, v in current["next"].items() if k != "result_id"} + current = page(item, token, **args) + assert seen == value and calls == 4 + + +def test_snapshot_is_immutable_and_bound_to_both_credential_and_principal(): + value = {"data": [1, 2]} + store = ResultStore() + token = store.put("list_keys", value, "alice-key", "user-1") + value["data"].append(3) + assert json.loads(store.get(token, "alice-key", "user-1").raw) == {"data": [1, 2]} + for credential, principal in [("bob-key", "user-1"), ("alice-key", "user-2")]: + with pytest.raises(ResultError, match="unavailable"): + store.get(token, credential, principal) + + +def test_expiration_capacity_and_clear_do_not_evict_unexpired_results(monkeypatch): + monkeypatch.setattr("litellm_admin_mcp.results.time.monotonic", lambda: 100) + store = ResultStore(ttl=10, max_bytes=100, max_per_owner=1) + token = store.put("get_user", {"a": 1}, "alice", "user-1") + assert store.put("get_user", {"b": 2}, "alice", "user-1") is None + assert store.put("get_user", {"a": "x" * 100}, "bob", "user-2") is None + assert store.get(token, "alice", "user-1") + monkeypatch.setattr("litellm_admin_mcp.results.time.monotonic", lambda: 111) + with pytest.raises(ResultError, match="expired"): + store.get(token, "alice", "user-1") + assert store._bytes == 0 + assert store.put("get_user", {"a": 1}, "alice", "user-1") + store.clear() + assert not store._items and store._bytes == 0 + + +@pytest.mark.parametrize("source,value", [ + ("create_key", {"key": "custom-key-without-prefix"}), + ("regenerate_key", {"key": "new"}), + ("regenerate_key_by_id", {"key": "new"}), + ("create_service_account_key", {"key": "new"}), + ("get_key", {"key": "custom-key-without-prefix", "info": {}}), + ("list_request_logs", {"history": "use sk-something here"}), + ("get_key", {"metadata": json.dumps({"key": "custom-secret"})}), + ("list_models", {"data": [{"model_info": {"key": "sk-sensitive"}}]}), + ("get_model", {"data": [{"model_info": {"key": "sk-sensitive"}}]}), + ("list_models", {"data": [{"model_info": {"key": json.dumps({"key": "custom-secret"})}}]}), + ("list_models", {"data": [{"model_info": {"key": "openai/example", "metadata": {"key": "custom-secret"}}}]}), + ("list_models", {"data": [{"key": "custom-secret"}]}), + ("list_models", {"model_info": {"key": "custom-secret"}}), + ("list_request_logs", {"data": [{"model_info": {"key": "custom-secret"}}]}), +]) +def test_credentials_are_not_retained_in_result_store(source, value): + store = ResultStore() + assert store.put(source, value, "alice", "user-1") is None + assert not store._items + + +@pytest.mark.parametrize("path", ["not-a-pointer", "/bad~2escape", "/array/-1", "/array/01", "/array/4", "/missing", "/scalar/x"]) +def test_invalid_paths_fail_explicitly(path): + with pytest.raises(ResultError): + select({"array": [1], "scalar": 0}, path) + + +def test_empty_values_and_pagination_edges(): + store = ResultStore() + for value in ([], {}, "", False, 0, None): + token = store.put("get_user", value, "alice", "user-1") + item = store.get(token, "alice", "user-1") + assert page(item, token)["complete"] + assert rebuild(item, token) == value + with pytest.raises(ResultError): + page(item, token, offset=1) diff --git a/tests/test_transports.py b/tests/test_transports.py index 2e3d718..0129518 100644 --- a/tests/test_transports.py +++ b/tests/test_transports.py @@ -6,6 +6,7 @@ from aiohttp import web from aiohttp.test_utils import TestServer import httpx2 +import pytest from mcp import ClientSession, StdioServerParameters from mcp.client.stdio import stdio_client from mcp.client.streamable_http import streamable_http_client @@ -13,7 +14,7 @@ from conftest import data from litellm_admin_mcp.config import Config -from litellm_admin_mcp.gateway import Gateway +from litellm_admin_mcp.gateway import Gateway, HELPERS from litellm_admin_mcp.server import create_http_app @@ -34,9 +35,10 @@ async def serve(app): await task -async def test_streamable_http_rejects_anonymous_and_keeps_callers_separate(stub): +@pytest.mark.parametrize("schema_mode", ["full", "discovery"]) +async def test_streamable_http_rejects_anonymous_and_keeps_callers_separate(stub, schema_mode): async with httpx2.AsyncClient(transport=httpx2.MockTransport(stub.handle)) as upstream: - app = create_http_app(Gateway(Config("https://gateway.example.com"), upstream)) + app = create_http_app(Gateway(Config("https://gateway.example.com", schema_mode=schema_mode), upstream)) async with serve(app) as base: async with httpx2.AsyncClient() as anonymous: assert (await anonymous.get(base + "/healthz")).status_code == 200 @@ -50,6 +52,12 @@ async def test_streamable_http_rejects_anonymous_and_keeps_callers_separate(stub assert "create_key" in {t.name for t in (await session.list_tools()).tools} result = data(await session.call_tool("create_key", {"body": {"key_alias": credential}})) assert result["key"] == "sk-created-for-[credential redacted]" + if credential == "alice-admin": + stub.keys = [{"metadata": "x" * 25000}] + saved = data(await session.call_tool("list_keys", {})) + assert data(await session.call_tool("read_admin_result", saved["full_read"])) == {"keys": stub.keys} + else: + assert (await session.call_tool("read_admin_result", saved["full_read"])).is_error assert [r.headers["Authorization"] for r in stub.writes] == ["Bearer alice-admin", "Bearer bob-admin"] @@ -71,8 +79,11 @@ async def handler(request): async with stdio_client(params) as (read, write): async with ClientSession(read, write) as session: await session.initialize() - assert {t.name for t in (await session.list_tools()).tools} == {"list_keys", "create_key"} + assert {t.name for t in (await session.list_tools()).tools} == {"list_keys", "create_key"} | HELPERS.keys() created = data(await session.call_tool("create_key", {"body": {"key_alias": "stdio"}})) assert created["key_alias"] == "stdio" + stub.keys = [{"metadata": "x" * 25000}] + saved = data(await session.call_tool("list_keys", {})) + assert data(await session.call_tool("read_admin_result", saved["full_read"])) == {"keys": stub.keys} assert len(stub.writes) == 1 assert "unrelated-slack-secret" not in json.dumps([dict(r.headers) for r in stub.requests])