From fb13413b81da24b3d0c525147cb3a26a0ca5ccb2 Mon Sep 17 00:00:00 2001 From: Andrej Simurka Date: Thu, 30 Apr 2026 16:07:20 +0200 Subject: [PATCH] Refactor of responses models dumping --- docs/openapi.json | 202 +++++++++--------- src/app/endpoints/responses.py | 7 +- src/models/api/requests/responses_openai.py | 5 +- src/models/common/responses/__init__.py | 4 + .../common/responses/responses_api_params.py | 27 +-- src/models/common/responses/types.py | 25 ++- src/utils/responses.py | 13 +- tests/unit/app/endpoints/test_responses.py | 32 ++- tests/unit/utils/test_responses.py | 7 +- tests/unit/utils/test_types.py | 22 +- 10 files changed, 177 insertions(+), 167 deletions(-) diff --git a/docs/openapi.json b/docs/openapi.json index 6dccf60cb..388dff21f 100644 --- a/docs/openapi.json +++ b/docs/openapi.json @@ -13506,6 +13506,105 @@ } ] }, + "InputToolMCP": { + "properties": { + "type": { + "type": "string", + "const": "mcp", + "title": "Type", + "default": "mcp" + }, + "server_label": { + "type": "string", + "title": "Server Label" + }, + "connector_id": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Connector Id" + }, + "server_url": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Server Url" + }, + "headers": { + "anyOf": [ + { + "additionalProperties": true, + "type": "object" + }, + { + "type": "null" + } + ], + "title": "Headers" + }, + "authorization": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Authorization" + }, + "require_approval": { + "anyOf": [ + { + "type": "string", + "const": "always" + }, + { + "type": "string", + "const": "never" + }, + { + "$ref": "#/components/schemas/ApprovalFilter" + } + ], + "title": "Require Approval", + "default": "never" + }, + "allowed_tools": { + "anyOf": [ + { + "items": { + "type": "string" + }, + "type": "array" + }, + { + "$ref": "#/components/schemas/AllowedToolsFilter" + }, + { + "type": "null" + } + ], + "title": "Allowed Tools" + } + }, + "type": "object", + "required": [ + "server_label" + ], + "title": "InputToolMCP", + "description": "MCP input tool with authorization included when serializing request bodies." + }, "InternalServerErrorResponse": { "properties": { "status_code": { @@ -15201,105 +15300,6 @@ "title": "OpenAIResponseInputToolFunction", "description": "Function tool configuration for OpenAI response inputs.\n\n:param type: Tool type identifier, always \"function\"\n:param name: Name of the function that can be called\n:param description: (Optional) Description of what the function does\n:param parameters: (Optional) JSON schema defining the function's parameters\n:param strict: (Optional) Whether to enforce strict parameter validation" }, - "OpenAIResponseInputToolMCP": { - "properties": { - "type": { - "type": "string", - "const": "mcp", - "title": "Type", - "default": "mcp" - }, - "server_label": { - "type": "string", - "title": "Server Label" - }, - "connector_id": { - "anyOf": [ - { - "type": "string" - }, - { - "type": "null" - } - ], - "title": "Connector Id" - }, - "server_url": { - "anyOf": [ - { - "type": "string" - }, - { - "type": "null" - } - ], - "title": "Server Url" - }, - "headers": { - "anyOf": [ - { - "additionalProperties": true, - "type": "object" - }, - { - "type": "null" - } - ], - "title": "Headers" - }, - "authorization": { - "anyOf": [ - { - "type": "string" - }, - { - "type": "null" - } - ], - "title": "Authorization" - }, - "require_approval": { - "anyOf": [ - { - "type": "string", - "const": "always" - }, - { - "type": "string", - "const": "never" - }, - { - "$ref": "#/components/schemas/ApprovalFilter" - } - ], - "title": "Require Approval", - "default": "never" - }, - "allowed_tools": { - "anyOf": [ - { - "items": { - "type": "string" - }, - "type": "array" - }, - { - "$ref": "#/components/schemas/AllowedToolsFilter" - }, - { - "type": "null" - } - ], - "title": "Allowed Tools" - } - }, - "type": "object", - "required": [ - "server_label" - ], - "title": "OpenAIResponseInputToolMCP", - "description": "Model Context Protocol (MCP) tool configuration for OpenAI response inputs.\n\n:param type: Tool type identifier, always \"mcp\"\n:param server_label: Label to identify this MCP server\n:param connector_id: (Optional) ID of the connector to use for this MCP server\n:param server_url: (Optional) URL endpoint of the MCP server\n:param headers: (Optional) HTTP headers to include when connecting to the server\n:param authorization: (Optional) OAuth access token for authenticating with the MCP server\n:param require_approval: Approval requirement for tool calls (\"always\", \"never\", or filter)\n:param allowed_tools: (Optional) Restriction on which tools can be used from this server" - }, "OpenAIResponseInputToolWebSearch": { "properties": { "type": { @@ -18082,7 +18082,7 @@ "$ref": "#/components/schemas/OpenAIResponseInputToolFunction" }, { - "$ref": "#/components/schemas/OpenAIResponseInputToolMCP" + "$ref": "#/components/schemas/InputToolMCP" } ], "discriminator": { @@ -18090,7 +18090,7 @@ "mapping": { "file_search": "#/components/schemas/OpenAIResponseInputToolFileSearch", "function": "#/components/schemas/OpenAIResponseInputToolFunction", - "mcp": "#/components/schemas/OpenAIResponseInputToolMCP", + "mcp": "#/components/schemas/InputToolMCP", "web_search": "#/components/schemas/OpenAIResponseInputToolWebSearch", "web_search_2025_08_26": "#/components/schemas/OpenAIResponseInputToolWebSearch", "web_search_preview": "#/components/schemas/OpenAIResponseInputToolWebSearch", diff --git a/src/app/endpoints/responses.py b/src/app/endpoints/responses.py index 2705a27dd..c52246591 100644 --- a/src/app/endpoints/responses.py +++ b/src/app/endpoints/responses.py @@ -455,12 +455,7 @@ async def responses_endpoint_handler( original_request.input, inline_rag_context.context_text ) - api_params = ResponsesApiParams.model_validate( - { - **updated_request.model_dump(exclude={"tools"}), - "tools": updated_request.tools, - } - ) + api_params = ResponsesApiParams.model_validate(updated_request.model_dump()) context = ResponsesContext( client=client, auth=auth, diff --git a/src/models/api/requests/responses_openai.py b/src/models/api/requests/responses_openai.py index 2b05f4767..dde597077 100644 --- a/src/models/api/requests/responses_openai.py +++ b/src/models/api/requests/responses_openai.py @@ -3,9 +3,6 @@ import json from typing import Any, Optional, Self -from llama_stack_api.openai_responses import ( - OpenAIResponseInputTool as InputTool, -) from llama_stack_api.openai_responses import ( OpenAIResponseInputToolChoice as ToolChoice, ) @@ -22,7 +19,7 @@ from constants import RESPONSES_REQUEST_MAX_SIZE from models.common.query import SolrVectorSearchRequest -from models.common.responses.types import IncludeParameter, ResponseInput +from models.common.responses.types import IncludeParameter, InputTool, ResponseInput from utils import suid diff --git a/src/models/common/responses/__init__.py b/src/models/common/responses/__init__.py index 6cdf7e5c7..9a7ad906f 100644 --- a/src/models/common/responses/__init__.py +++ b/src/models/common/responses/__init__.py @@ -7,6 +7,8 @@ ) from models.common.responses.types import ( IncludeParameter, + InputTool, + InputToolMCP, ResponseInput, ResponseItem, ) @@ -15,6 +17,8 @@ "ResponseInput", "ResponseItem", "IncludeParameter", + "InputTool", + "InputToolMCP", "ResponsesApiParams", "ResponsesContext", "ResponsesConversationContext", diff --git a/src/models/common/responses/responses_api_params.py b/src/models/common/responses/responses_api_params.py index acb219c89..6abf9a618 100644 --- a/src/models/common/responses/responses_api_params.py +++ b/src/models/common/responses/responses_api_params.py @@ -3,9 +3,6 @@ from collections.abc import Mapping from typing import Any, Final, Optional -from llama_stack_api.openai_responses import ( - OpenAIResponseInputTool as InputTool, -) from llama_stack_api.openai_responses import ( OpenAIResponseInputToolChoice as ToolChoice, ) @@ -23,7 +20,7 @@ ) from pydantic import BaseModel, Field -from models.common.responses.types import IncludeParameter, ResponseInput +from models.common.responses.types import IncludeParameter, InputTool, ResponseInput from utils.tool_formatter import translate_vector_store_ids_to_user_facing # Attribute names that are echoed back in the response. @@ -126,28 +123,16 @@ class ResponsesApiParams(BaseModel): ) def model_dump(self, *args: Any, **kwargs: Any) -> dict[str, Any]: - """Serialize params, re-injecting MCP authorization stripped by exclude=True. + """Serialize to a request body dict. - llama-stack-api marks ``InputToolMCP.authorization`` with - ``Field(exclude=True)`` to prevent token leakage in API responses. - The base ``model_dump()`` therefore strips the field, but we need it - in the request payload so llama-stack server can authenticate with - MCP servers. See LCORE-1414 / GitHub issue #1269. + Omits conversation when previous_response_id is set. + + Returns: + Serializable dict for the Responses API request body. """ result = super().model_dump(*args, **kwargs) - # Only one context option is allowed, previous_response_id has priority - # Turn is added to conversation manually if previous_response_id is used if self.previous_response_id: result.pop("conversation", None) - dumped_tools = result.get("tools") - if not self.tools or not isinstance(dumped_tools, list): - return result - if len(dumped_tools) != len(self.tools): - return result - for tool, dumped_tool in zip(self.tools, dumped_tools): - authorization = getattr(tool, "authorization", None) - if authorization is not None and isinstance(dumped_tool, dict): - dumped_tool["authorization"] = authorization return result def echoed_params(self, rag_id_mapping: Mapping[str, str]) -> dict[str, Any]: diff --git a/src/models/common/responses/types.py b/src/models/common/responses/types.py index 992d5a4df..a9a5710ed 100644 --- a/src/models/common/responses/types.py +++ b/src/models/common/responses/types.py @@ -1,10 +1,20 @@ """Type aliases for OpenAI-compatible Responses API input shapes.""" -from typing import Literal +from typing import Annotated, Literal from llama_stack_api.openai_responses import ( OpenAIResponseInputFunctionToolCallOutput as FunctionToolCallOutput, ) +from llama_stack_api.openai_responses import ( + OpenAIResponseInputToolFileSearch as InputToolFileSearch, +) +from llama_stack_api.openai_responses import ( + OpenAIResponseInputToolFunction as InputToolFunction, +) +from llama_stack_api.openai_responses import OpenAIResponseInputToolMCP +from llama_stack_api.openai_responses import ( + OpenAIResponseInputToolWebSearch as InputToolWebSearch, +) from llama_stack_api.openai_responses import ( OpenAIResponseMCPApprovalRequest as McpApprovalRequest, ) @@ -29,6 +39,19 @@ from llama_stack_api.openai_responses import ( OpenAIResponseOutputMessageWebSearchToolCall as WebSearchToolCall, ) +from pydantic import Field + + +class InputToolMCP(OpenAIResponseInputToolMCP): + """MCP input tool with authorization included when serializing request bodies.""" + + authorization: str | None = None + + +InputTool = Annotated[ + InputToolWebSearch | InputToolFileSearch | InputToolFunction | InputToolMCP, + Field(discriminator="type"), +] type IncludeParameter = Literal[ "web_search_call.action.sources", diff --git a/src/utils/responses.py b/src/utils/responses.py index 0ed7477bc..385d5dad2 100644 --- a/src/utils/responses.py +++ b/src/utils/responses.py @@ -20,9 +20,6 @@ from llama_stack_api.openai_responses import ( OpenAIResponseInputMessageContentText as InputTextPart, ) -from llama_stack_api.openai_responses import ( - OpenAIResponseInputTool as InputTool, -) from llama_stack_api.openai_responses import ( OpenAIResponseInputToolChoice as ToolChoice, ) @@ -35,9 +32,6 @@ from llama_stack_api.openai_responses import ( OpenAIResponseInputToolFileSearch as InputToolFileSearch, ) -from llama_stack_api.openai_responses import ( - OpenAIResponseInputToolMCP as InputToolMCP, -) from llama_stack_api.openai_responses import ( OpenAIResponseMCPApprovalRequest as MCPApprovalRequest, ) @@ -99,7 +93,12 @@ ServiceUnavailableResponse, ) from models.common.responses.responses_api_params import ResponsesApiParams -from models.common.responses.types import ResponseInput, ResponseItem +from models.common.responses.types import ( + InputTool, + InputToolMCP, + ResponseInput, + ResponseItem, +) from models.common.turn_summary import ( RAGChunk, ReferencedDocument, diff --git a/tests/unit/app/endpoints/test_responses.py b/tests/unit/app/endpoints/test_responses.py index b5d2e5bd9..ed95708ab 100644 --- a/tests/unit/app/endpoints/test_responses.py +++ b/tests/unit/app/endpoints/test_responses.py @@ -13,7 +13,9 @@ from llama_stack_api.openai_responses import ( OpenAIResponseInputToolChoiceMode as ToolChoiceMode, ) -from llama_stack_api.openai_responses import OpenAIResponseMessage +from llama_stack_api.openai_responses import ( + OpenAIResponseMessage, +) from llama_stack_client import APIConnectionError, APIStatusError, AsyncLlamaStackClient from pytest_mock import MockerFixture @@ -37,6 +39,7 @@ from models.common.responses.responses_conversation_context import ( ResponsesConversationContext, ) +from models.common.responses.types import InputToolMCP from models.common.turn_summary import RAGContext, TurnSummary from models.config import Action, ModelContextProtocolServer from models.database.conversations import UserConversation @@ -73,12 +76,7 @@ def build_api_params_and_context( # pylint: disable=too-many-arguments user_agent: Optional[str] = None, ) -> tuple[ResponsesApiParams, ResponsesContext]: """Build api_params/context for direct helper invocation tests.""" - api_params = ResponsesApiParams.model_validate( - { - **updated_request.model_dump(exclude={"tools"}), - "tools": updated_request.tools, - } - ) + api_params = ResponsesApiParams.model_validate(updated_request.model_dump()) context = ResponsesContext.model_construct( client=client, auth=auth, @@ -96,6 +94,26 @@ def build_api_params_and_context( # pylint: disable=too-many-arguments return api_params, context +def test_responses_api_params_preserves_mcp_authorization() -> None: + """After model_validate, MCP tool authorization from model_dump is kept on api_params.tools.""" + token = "secret-token" + req = ResponsesRequest( + input="x", + model=MODEL, + conversation=VALID_CONV_ID, + tools=[ + InputToolMCP( + server_label="alpha", + server_url="http://alpha", + require_approval="never", + authorization=token, + ) + ], + ) + api = ResponsesApiParams.model_validate(req.model_dump()) + assert api.tools is not None and api.tools[0].authorization == token + + def _patch_base(mocker: MockerFixture, config: AppConfig) -> None: """Patch configuration and mandatory checks for responses endpoint.""" mocker.patch(f"{MODULE}.configuration", config) diff --git a/tests/unit/utils/test_responses.py b/tests/unit/utils/test_responses.py index f338f5401..21dd76b2e 100644 --- a/tests/unit/utils/test_responses.py +++ b/tests/unit/utils/test_responses.py @@ -12,9 +12,6 @@ AllowedToolsFilter, OpenAIResponseInputToolChoiceAllowedTools, ) -from llama_stack_api.openai_responses import ( - OpenAIResponseInputTool as InputTool, -) from llama_stack_api.openai_responses import ( OpenAIResponseInputToolChoiceFileSearch as ToolChoiceFileSearch, ) @@ -27,9 +24,6 @@ from llama_stack_api.openai_responses import ( OpenAIResponseInputToolFunction as InputToolFunction, ) -from llama_stack_api.openai_responses import ( - OpenAIResponseInputToolMCP as InputToolMCP, -) from llama_stack_api.openai_responses import ( OpenAIResponseInputToolWebSearch as InputToolWebSearch, ) @@ -60,6 +54,7 @@ import constants from models.api.requests import QueryRequest +from models.common.responses.types import InputTool, InputToolMCP from models.config import ByokRag, ModelContextProtocolServer from utils.responses import ( _build_chunk_attributes, diff --git a/tests/unit/utils/test_types.py b/tests/unit/utils/test_types.py index 8447054f0..eb20f13ae 100644 --- a/tests/unit/utils/test_types.py +++ b/tests/unit/utils/test_types.py @@ -2,15 +2,10 @@ import pytest from llama_stack_api import URL, ImageContentItem, TextContentItem, _URLOrData -from llama_stack_api.openai_responses import ( - OpenAIResponseInputToolFileSearch as InputToolFileSearch, -) -from llama_stack_api.openai_responses import ( - OpenAIResponseInputToolMCP as InputToolMCP, -) from pydantic import AnyUrl, ValidationError from models.common.responses.responses_api_params import ResponsesApiParams +from models.common.responses.types import InputToolFileSearch, InputToolMCP from models.common.turn_summary import ( ReferencedDocument, ToolCallSummary, @@ -148,11 +143,10 @@ def test_constructor_partial_fields(self) -> None: class TestResponsesApiParamsModelDump: - """Tests for ResponsesApiParams.model_dump() MCP authorization fix. + """Tests for ResponsesApiParams.model_dump() MCP authorization serialization. - Regression tests for LCORE-1414 / GitHub issue #1269: llama-stack-api's - InputToolMCP.authorization has Field(exclude=True), causing the base - model_dump() to silently strip authorization tokens. + Regression tests for LCORE-1414 / GitHub issue #1269: MCP authorization must + survive model_dump() when forwarding tools to Llama Stack. """ def _make_params(self, tools: list) -> ResponsesApiParams: @@ -175,7 +169,7 @@ def test_mcp_authorization_survives_model_dump(self) -> None: authorization="my-secret-token", ) assert tool.authorization == "my-secret-token" - assert "authorization" not in tool.model_dump() + assert tool.model_dump()["authorization"] == "my-secret-token" params = self._make_params([tool]) dumped = params.model_dump(exclude_none=True) @@ -240,8 +234,8 @@ def test_multiple_mcp_tools_each_preserves_authorization(self) -> None: assert dumped["tools"][0]["authorization"] == "token-a" assert dumped["tools"][1]["authorization"] == "token-b" - def test_exclude_changing_tool_list_shape_skips_reinjection(self) -> None: - """Test that exclude removing tool indices does not mis-assign authorization.""" + def test_partial_tool_dump_reinjects_auth_by_server_label(self) -> None: + """When exclude drops some tools, remaining MCP rows still get auth by label.""" tool_a = InputToolMCP( server_label="server-a", server_url="http://a:3000", @@ -258,7 +252,7 @@ def test_exclude_changing_tool_list_shape_skips_reinjection(self) -> None: dumped = params.model_dump(exclude={"tools": {0}}) assert len(dumped["tools"]) == 1 assert dumped["tools"][0]["server_label"] == "server-b" - assert "authorization" not in dumped["tools"][0] + assert dumped["tools"][0]["authorization"] == "token-b" def test_no_tools_does_not_error(self) -> None: """Test that model_dump() works when tools is None."""