diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260209085821_add_verificationtoken_indexes/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260209085821_add_verificationtoken_indexes/migration.sql new file mode 100644 index 000000000000..572eea9b5295 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260209085821_add_verificationtoken_indexes/migration.sql @@ -0,0 +1,8 @@ +-- CreateIndex +CREATE INDEX "LiteLLM_VerificationToken_user_id_team_id_idx" ON "LiteLLM_VerificationToken"("user_id", "team_id"); + +-- CreateIndex +CREATE INDEX "LiteLLM_VerificationToken_team_id_idx" ON "LiteLLM_VerificationToken"("team_id"); + +-- CreateIndex +CREATE INDEX "LiteLLM_VerificationToken_budget_reset_at_expires_idx" ON "LiteLLM_VerificationToken"("budget_reset_at", "expires"); diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index c2a599c178fc..b1ca1f71c9e6 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -310,6 +310,16 @@ model LiteLLM_VerificationToken { litellm_budget_table LiteLLM_BudgetTable? @relation(fields: [budget_id], references: [budget_id]) litellm_organization_table LiteLLM_OrganizationTable? @relation(fields: [organization_id], references: [organization_id]) object_permission LiteLLM_ObjectPermissionTable? @relation(fields: [object_permission_id], references: [object_permission_id]) + + // SELECT COUNT(*) FROM (SELECT "public"."LiteLLM_VerificationToken"."token" FROM "public"."LiteLLM_VerificationToken" WHERE ("public"."LiteLLM_VerificationToken"."user_id" = $1 AND ("public"."LiteLLM_VerificationToken"."team_id" IS NULL OR "public"."LiteLLM_VerificationToken"."team_id" <> $2)) OFFSET $3 ) AS "sub" + // SELECT ... FROM "public"."LiteLLM_VerificationToken" WHERE "public"."LiteLLM_VerificationToken"."user_id" = $1 OFFSET $2 + @@index([user_id, team_id]) + + // SELECT ... FROM "public"."LiteLLM_VerificationToken" WHERE "public"."LiteLLM_VerificationToken"."team_id" = $1 OFFSET $2 + @@index([team_id]) + + // SELECT ... FROM "public"."LiteLLM_VerificationToken" WHERE (("public"."LiteLLM_VerificationToken"."expires" IS NULL OR "public"."LiteLLM_VerificationToken"."expires" > $1) AND "public"."LiteLLM_VerificationToken"."budget_reset_at" < $2) OFFSET $3 + @@index([budget_reset_at, expires]) } // Audit table for deleted keys - preserves spend and key information for historical tracking diff --git a/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py b/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py index b635f15ed09f..deb0b4f95492 100644 --- a/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py +++ b/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py @@ -3,6 +3,8 @@ """ import json +import asyncio +import os from pathlib import PurePosixPath from typing import Any, Dict, Optional from urllib.parse import quote @@ -45,8 +47,36 @@ def _sanitize_path_parameter_value(param_value: Any, param_name: str) -> str: def load_openapi_spec(filepath: str) -> Dict[str, Any]: - """Load OpenAPI specification from JSON file.""" - with open(filepath, "r") as f: + """ + Sync wrapper. For URL specs, use the shared/custom MCP httpx client. + """ + try: + # If we're already inside an event loop, prefer the async function. + asyncio.get_running_loop() + raise RuntimeError( + "load_openapi_spec() was called from within a running event loop. " + "Use 'await load_openapi_spec_async(...)' instead." + ) + except RuntimeError as e: + # "no running event loop" is fine; other RuntimeErrors we re-raise + if "no running event loop" not in str(e).lower(): + raise + return asyncio.run(load_openapi_spec_async(filepath)) + +async def load_openapi_spec_async(filepath: str) -> Dict[str, Any]: + if filepath.startswith("http://") or filepath.startswith("https://"): + client = get_async_httpx_client(llm_provider=httpxSpecialProvider.MCP) + # NOTE: do not close shared client if get_async_httpx_client returns a shared singleton. + # If it returns a new client each time, consider wrapping it in an async context manager. + r = await client.get(filepath) + r.raise_for_status() + return r.json() + + # fallback: local file + # Local filesystem path + if not os.path.exists(filepath): + raise FileNotFoundError(f"OpenAPI spec not found at {filepath}") + with open(filepath, "r", encoding="utf-8") as f: return json.load(f) diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 279946f78def..1750efed92c3 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -308,6 +308,16 @@ model LiteLLM_VerificationToken { litellm_budget_table LiteLLM_BudgetTable? @relation(fields: [budget_id], references: [budget_id]) litellm_organization_table LiteLLM_OrganizationTable? @relation(fields: [organization_id], references: [organization_id]) object_permission LiteLLM_ObjectPermissionTable? @relation(fields: [object_permission_id], references: [object_permission_id]) + + // SELECT COUNT(*) FROM (SELECT "public"."LiteLLM_VerificationToken"."token" FROM "public"."LiteLLM_VerificationToken" WHERE ("public"."LiteLLM_VerificationToken"."user_id" = $1 AND ("public"."LiteLLM_VerificationToken"."team_id" IS NULL OR "public"."LiteLLM_VerificationToken"."team_id" <> $2)) OFFSET $3 ) AS "sub" + // SELECT ... FROM "public"."LiteLLM_VerificationToken" WHERE "public"."LiteLLM_VerificationToken"."user_id" = $1 OFFSET $2 + @@index([user_id, team_id]) + + // SELECT ... FROM "public"."LiteLLM_VerificationToken" WHERE "public"."LiteLLM_VerificationToken"."team_id" = $1 OFFSET $2 + @@index([team_id]) + + // SELECT ... FROM "public"."LiteLLM_VerificationToken" WHERE (("public"."LiteLLM_VerificationToken"."expires" IS NULL OR "public"."LiteLLM_VerificationToken"."expires" > $1) AND "public"."LiteLLM_VerificationToken"."budget_reset_at" < $2) OFFSET $3 + @@index([budget_reset_at, expires]) } // Audit table for deleted keys - preserves spend and key information for historical tracking diff --git a/litellm/responses/litellm_completion_transformation/streaming_iterator.py b/litellm/responses/litellm_completion_transformation/streaming_iterator.py index 867c18b6dd4d..5c05526442d2 100644 --- a/litellm/responses/litellm_completion_transformation/streaming_iterator.py +++ b/litellm/responses/litellm_completion_transformation/streaming_iterator.py @@ -88,6 +88,8 @@ def __init__( self._pending_tool_events: List[BaseLiteLLMOpenAIResponseObject] = [] self._tool_output_index_by_call_id: dict[str, int] = {} self._tool_args_by_call_id: dict[str, str] = {} + self._tool_call_id_by_index: dict[int, str] = {} + self._ambiguous_tool_call_indexes: set[int] = set() self._next_tool_output_index: int = 1 # output_index=0 reserved for the message item self._final_tool_events_queued: bool = False self._sequence_number: int = 0 @@ -111,6 +113,19 @@ def _get_or_assign_tool_output_index(self, call_id: str) -> int: self._tool_output_index_by_call_id[call_id] = idx return idx + def _normalize_tool_call_index(self, tool_call: object) -> Optional[int]: + idx_raw = ( + tool_call.get("index") + if isinstance(tool_call, dict) + else getattr(tool_call, "index", None) + ) + if idx_raw is None: + return None + try: + return int(idx_raw) + except (TypeError, ValueError): + return None + def _is_reasoning_end(self, chunk): delta = chunk.choices[0].delta @@ -143,10 +158,28 @@ def _queue_tool_call_delta_events(self, tool_calls: object) -> None: return for tc in tool_calls: + tc_index = self._normalize_tool_call_index(tc) call_id_raw = tc.get("id") if isinstance(tc, dict) else getattr(tc, "id", None) - if not call_id_raw: + call_id = "" + + if call_id_raw: + call_id = str(call_id_raw) + if tc_index is not None: + existing_call_id = self._tool_call_id_by_index.get(tc_index) + if existing_call_id is not None and existing_call_id != call_id: + # Reusing the same index for multiple call_ids is ambiguous for id-less deltas. + # Guard against silent misrouting by disabling index fallback for this index. + self._ambiguous_tool_call_indexes.add(tc_index) + self._tool_call_id_by_index[tc_index] = call_id + elif tc_index is not None: + if tc_index in self._ambiguous_tool_call_indexes: + continue + mapped_call_id = self._tool_call_id_by_index.get(tc_index) + if mapped_call_id: + call_id = mapped_call_id + + if not call_id: continue - call_id = str(call_id_raw) fn = tc.get("function") if isinstance(tc, dict) else getattr(tc, "function", None) fn_name = "" diff --git a/schema.prisma b/schema.prisma index ecf4e06ef6e8..9a87a491cf72 100644 --- a/schema.prisma +++ b/schema.prisma @@ -310,6 +310,16 @@ model LiteLLM_VerificationToken { litellm_budget_table LiteLLM_BudgetTable? @relation(fields: [budget_id], references: [budget_id]) litellm_organization_table LiteLLM_OrganizationTable? @relation(fields: [organization_id], references: [organization_id]) object_permission LiteLLM_ObjectPermissionTable? @relation(fields: [object_permission_id], references: [object_permission_id]) + + // SELECT COUNT(*) FROM (SELECT "public"."LiteLLM_VerificationToken"."token" FROM "public"."LiteLLM_VerificationToken" WHERE ("public"."LiteLLM_VerificationToken"."user_id" = $1 AND ("public"."LiteLLM_VerificationToken"."team_id" IS NULL OR "public"."LiteLLM_VerificationToken"."team_id" <> $2)) OFFSET $3 ) AS "sub" + // SELECT ... FROM "public"."LiteLLM_VerificationToken" WHERE "public"."LiteLLM_VerificationToken"."user_id" = $1 OFFSET $2 + @@index([user_id, team_id]) + + // SELECT ... FROM "public"."LiteLLM_VerificationToken" WHERE "public"."LiteLLM_VerificationToken"."team_id" = $1 OFFSET $2 + @@index([team_id]) + + // SELECT ... FROM "public"."LiteLLM_VerificationToken" WHERE (("public"."LiteLLM_VerificationToken"."expires" IS NULL OR "public"."LiteLLM_VerificationToken"."expires" > $1) AND "public"."LiteLLM_VerificationToken"."budget_reset_at" < $2) OFFSET $3 + @@index([budget_reset_at, expires]) } // Audit table for deleted keys - preserves spend and key information for historical tracking diff --git a/tests/mcp_tests/test_openapi_spec_path_url.py b/tests/mcp_tests/test_openapi_spec_path_url.py new file mode 100644 index 000000000000..03e9db949679 --- /dev/null +++ b/tests/mcp_tests/test_openapi_spec_path_url.py @@ -0,0 +1,92 @@ +from __future__ import annotations + +from typing import Any, Dict + +import httpx +import pytest + +from litellm.proxy._experimental.mcp_server import openapi_to_mcp_generator as gen + + +class _FakeAsyncHTTPHandler: + """ + Minimal stand-in for the object returned by get_async_httpx_client(). + openapi_to_mcp_generator.load_openapi_spec_async() calls: + + client = get_async_httpx_client(...) + r = await client.get(url, timeout=30.0) + + So we must implement async get(). + """ + + def __init__(self, response: httpx.Response, expected_url: str): + self._response = response + self._expected_url = expected_url + self.calls = 0 + + async def get(self, request_url: str, timeout: float = 30.0): + self.calls += 1 + assert request_url == self._expected_url + assert timeout == 30.0 + return self._response + + +def test_load_openapi_spec_supports_http_url(monkeypatch: pytest.MonkeyPatch) -> None: + url = "http://example.local/openapi.json" + expected: Dict[str, Any] = { + "openapi": "3.0.0", + "info": {"title": "Test API", "version": "1.0.0"}, + "paths": {}, + } + + # httpx.Response must include a Request for raise_for_status() to work. + req = httpx.Request("GET", url) + resp = httpx.Response(status_code=200, json=expected, request=req) + + calls = {"get_async_httpx_client": 0} + handler_holder: Dict[str, Any] = {} + + def fake_get_async_httpx_client(*args, **kwargs): + calls["get_async_httpx_client"] += 1 + h = _FakeAsyncHTTPHandler(resp, expected_url=url) + handler_holder["handler"] = h + return h + + # Ensure shared/custom client path is used + monkeypatch.setattr(gen, "get_async_httpx_client", fake_get_async_httpx_client) + + # Fail loudly if someone reintroduces direct httpx.get() + def boom(*args, **kwargs): + raise AssertionError("Direct httpx.get() must not be used for URL spec loading") + + monkeypatch.setattr(httpx, "get", boom) + + spec = gen.load_openapi_spec(url) + + assert spec == expected + assert calls["get_async_httpx_client"] == 1 + assert handler_holder["handler"].calls == 1 + + +def test_load_openapi_spec_supports_local_file_path(tmp_path, monkeypatch: pytest.MonkeyPatch) -> None: + expected: Dict[str, Any] = { + "openapi": "3.0.0", + "info": {"title": "Local API", "version": "1.0.0"}, + "paths": {}, + } + + p = tmp_path / "openapi.json" + p.write_text( + '{"openapi":"3.0.0","info":{"title":"Local API","version":"1.0.0"},"paths":{}}', + encoding="utf-8", + ) + + # For local files, shared client must NOT be used. + def boom_client(*args, **kwargs): + raise AssertionError("get_async_httpx_client() must not be called for local file paths") + + monkeypatch.setattr(gen, "get_async_httpx_client", boom_client) + + spec = gen.load_openapi_spec(str(p)) + assert spec == expected + diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_tool_call_streaming_transformation.py b/tests/test_litellm/responses/litellm_completion_transformation/test_tool_call_streaming_transformation.py index 8d324bea6112..071eefaef477 100644 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_tool_call_streaming_transformation.py +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_tool_call_streaming_transformation.py @@ -229,3 +229,164 @@ def test_tool_call_arguments_are_chunked_to_match_openai_behavior(): assert sequence_numbers == sorted(sequence_numbers) assert len(set(sequence_numbers)) == len(sequence_numbers) # All unique + +def test_tool_call_delta_without_id_uses_index_mapping(): + iterator = LiteLLMCompletionStreamingIterator( + model="test-model", + litellm_custom_stream_wrapper=AsyncMock(), + request_input="Test input", + responses_api_request={}, + ) + + chunks = [ + [ + { + "index": 0, + "id": "call_abc123", + "type": "function", + "function": {"name": "get_weather", "arguments": '{"lo'}, + } + ], + [{"index": 0, "type": "function", "function": {"arguments": 'cation":'}}], + [{"index": 0, "type": "function", "function": {"arguments": ' "New'}}], + [{"index": 0, "type": "function", "function": {"arguments": ' York"}'}}], + ] + + for tool_calls in chunks: + iterator._queue_tool_call_delta_events(tool_calls) + + all_events = [] + while iterator._pending_tool_events: + all_events.append(iterator._pending_tool_events.pop(0)) + + delta_events = [ + evt + for evt in all_events + if evt.type == ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DELTA + ] + streamed_arguments = "".join(evt.delta for evt in delta_events) + + assert streamed_arguments == '{"location": "New York"}' + + output_item_added_events = [ + evt + for evt in all_events + if evt.type == ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED + ] + assert len(output_item_added_events) == 1 + assert output_item_added_events[0].item.id == "call_abc123" + + +def test_parallel_tool_calls_without_ids_use_index_mapping(): + iterator = LiteLLMCompletionStreamingIterator( + model="test-model", + litellm_custom_stream_wrapper=AsyncMock(), + request_input="Test input", + responses_api_request={}, + ) + + iterator._queue_tool_call_delta_events( + [ + { + "index": 0, + "id": "call_a", + "type": "function", + "function": {"name": "tool_a", "arguments": '{"x":'}, + }, + { + "index": 1, + "id": "call_b", + "type": "function", + "function": {"name": "tool_b", "arguments": '{"y":'}, + }, + ] + ) + iterator._queue_tool_call_delta_events( + [ + {"index": 0, "type": "function", "function": {"arguments": "1}"}}, + {"index": 1, "type": "function", "function": {"arguments": "2}"}}, + ] + ) + + all_events = [] + while iterator._pending_tool_events: + all_events.append(iterator._pending_tool_events.pop(0)) + + output_item_added_events = [ + evt + for evt in all_events + if evt.type == ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED + ] + assert len(output_item_added_events) == 2 + + delta_events = [ + evt + for evt in all_events + if evt.type == ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DELTA + ] + arguments_by_call_id = {} + for evt in delta_events: + arguments_by_call_id.setdefault(evt.item_id, "") + arguments_by_call_id[evt.item_id] += evt.delta + + assert arguments_by_call_id["call_a"] == '{"x":1}' + assert arguments_by_call_id["call_b"] == '{"y":2}' + + +def test_reused_index_with_new_call_id_marks_fallback_ambiguous(): + iterator = LiteLLMCompletionStreamingIterator( + model="test-model", + litellm_custom_stream_wrapper=AsyncMock(), + request_input="Test input", + responses_api_request={}, + ) + + iterator._queue_tool_call_delta_events( + [ + { + "index": 0, + "id": "call_a", + "type": "function", + "function": {"name": "tool_a", "arguments": '{"a":'}, + } + ] + ) + iterator._queue_tool_call_delta_events( + [ + { + "index": 0, + "id": "call_b", + "type": "function", + "function": {"name": "tool_b", "arguments": '{"b":'}, + } + ] + ) + # Ambiguous chunk: index reused and id missing. We should skip fallback rather than misroute. + iterator._queue_tool_call_delta_events( + [ + { + "index": 0, + "type": "function", + "function": {"arguments": "1}"}, + } + ] + ) + + all_events = [] + while iterator._pending_tool_events: + all_events.append(iterator._pending_tool_events.pop(0)) + + delta_events = [ + evt + for evt in all_events + if evt.type == ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DELTA + ] + arguments_by_call_id = {} + for evt in delta_events: + arguments_by_call_id.setdefault(evt.item_id, "") + arguments_by_call_id[evt.item_id] += evt.delta + + assert arguments_by_call_id["call_a"] == '{"a":' + assert arguments_by_call_id["call_b"] == '{"b":' + assert arguments_by_call_id["call_a"] != '{"a":1}' + assert arguments_by_call_id["call_b"] != '{"b":1}'