Skip to content

Commit b0fe07d

Browse files
authored
Merge pull request #1275 from max-svistunov/lcore-1414-fix-mcp-authorization-serialization
LCORE-1414 Fix MCP authorization loss during model_dump() serialization
2 parents fa8e7e9 + 41a7a58 commit b0fe07d

9 files changed

Lines changed: 242 additions & 0 deletions

File tree

docker-compose.yaml

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -81,6 +81,7 @@ services:
8181
- "8080:8080"
8282
volumes:
8383
- ./lightspeed-stack.yaml:/app-root/lightspeed-stack.yaml:z
84+
- ./tests/e2e/secrets/mcp-token:/tmp/mcp-secret-token:ro
8485
environment:
8586
- OPENAI_API_KEY=${OPENAI_API_KEY}
8687
# Azure Entra ID credentials (AZURE_API_KEY is obtained dynamically)

src/utils/types.py

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -220,6 +220,27 @@ class ResponsesApiParams(BaseModel):
220220
description="Extra HTTP headers to send with the request (e.g. x-llamastack-provider-data)",
221221
)
222222

223+
def model_dump(self, *args: Any, **kwargs: Any) -> dict[str, Any]:
224+
"""Serialize params, re-injecting MCP authorization stripped by exclude=True.
225+
226+
llama-stack-api marks ``InputToolMCP.authorization`` with
227+
``Field(exclude=True)`` to prevent token leakage in API responses.
228+
The base ``model_dump()`` therefore strips the field, but we need it
229+
in the request payload so llama-stack server can authenticate with
230+
MCP servers. See LCORE-1414 / GitHub issue #1269.
231+
"""
232+
result = super().model_dump(*args, **kwargs)
233+
dumped_tools = result.get("tools")
234+
if not self.tools or not isinstance(dumped_tools, list):
235+
return result
236+
if len(dumped_tools) != len(self.tools):
237+
return result
238+
for tool, dumped_tool in zip(self.tools, dumped_tools):
239+
authorization = getattr(tool, "authorization", None)
240+
if authorization is not None and isinstance(dumped_tool, dict):
241+
dumped_tool["authorization"] = authorization
242+
return result
243+
223244

224245
class ToolCallSummary(BaseModel):
225246
"""Model representing a tool call made during response generation (for tool_calls list)."""
Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,25 @@
1+
name: Lightspeed Core Service (LCS)
2+
service:
3+
host: 0.0.0.0
4+
port: 8080
5+
auth_enabled: false
6+
workers: 1
7+
color_log: true
8+
access_log: true
9+
llama_stack:
10+
# Library mode - embeds llama-stack as library
11+
use_as_library_client: true
12+
library_client_config_path: run.yaml
13+
user_data_collection:
14+
feedback_enabled: true
15+
feedback_storage: "/tmp/data/feedback"
16+
transcripts_enabled: true
17+
transcripts_storage: "/tmp/data/transcripts"
18+
authentication:
19+
module: "noop"
20+
mcp_servers:
21+
- name: "mcp-file-auth"
22+
provider_id: "model-context-protocol"
23+
url: "http://mock-mcp:3001"
24+
authorization_headers:
25+
Authorization: "/tmp/mcp-secret-token"
Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,26 @@
1+
name: Lightspeed Core Service (LCS)
2+
service:
3+
host: 0.0.0.0
4+
port: 8080
5+
auth_enabled: false
6+
workers: 1
7+
color_log: true
8+
access_log: true
9+
llama_stack:
10+
# Server mode - connects to separate llama-stack service
11+
use_as_library_client: false
12+
url: http://llama-stack:8321
13+
api_key: xyzzy
14+
user_data_collection:
15+
feedback_enabled: true
16+
feedback_storage: "/tmp/data/feedback"
17+
transcripts_enabled: true
18+
transcripts_storage: "/tmp/data/transcripts"
19+
authentication:
20+
module: "noop"
21+
mcp_servers:
22+
- name: "mcp-file-auth"
23+
provider_id: "model-context-protocol"
24+
url: "http://mock-mcp:3001"
25+
authorization_headers:
26+
Authorization: "/tmp/mcp-secret-token"

tests/e2e/features/environment.py

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -53,6 +53,10 @@
5353
"tests/e2e/configuration/{mode_dir}/lightspeed-stack-auth-rh-identity.yaml",
5454
"tests/e2e-prow/rhoai/configs/lightspeed-stack-auth-rh-identity.yaml",
5555
),
56+
"mcp-file-auth": (
57+
"tests/e2e/configuration/{mode_dir}/lightspeed-stack-mcp-file-auth.yaml",
58+
"tests/e2e-prow/rhoai/configs/lightspeed-stack-mcp-file-auth.yaml",
59+
),
5660
}
5761

5862

@@ -383,6 +387,12 @@ def before_feature(context: Context, feature: Feature) -> None:
383387
switch_config(context.feature_config)
384388
restart_container("lightspeed-stack")
385389

390+
if "MCPFileAuth" in feature.tags:
391+
context.feature_config = _get_config_path("mcp-file-auth", mode_dir)
392+
context.default_config_backup = create_config_backup("lightspeed-stack.yaml")
393+
switch_config(context.feature_config)
394+
restart_container("lightspeed-stack")
395+
386396

387397
def after_feature(context: Context, feature: Feature) -> None:
388398
"""Run after each feature file is exercised.
@@ -415,3 +425,8 @@ def after_feature(context: Context, feature: Feature) -> None:
415425
switch_config(context.default_config_backup)
416426
restart_container("lightspeed-stack")
417427
remove_config_backup(context.default_config_backup)
428+
429+
if "MCPFileAuth" in feature.tags:
430+
switch_config(context.default_config_backup)
431+
restart_container("lightspeed-stack")
432+
remove_config_backup(context.default_config_backup)
Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,20 @@
1+
@MCPFileAuth
2+
Feature: MCP file-based authorization tests
3+
4+
Regression tests for LCORE-1414: MCP authorization tokens configured via
5+
file-based authorization_headers must survive model_dump() serialization
6+
and reach the MCP server as a valid Bearer token.
7+
8+
Background:
9+
Given The service is started locally
10+
And REST API service prefix is /v1
11+
12+
@skip-in-library-mode
13+
Scenario: Query succeeds with file-based MCP authorization
14+
Given The system is in default state
15+
When I use "query" to ask question
16+
"""
17+
{"query": "Use the mock_tool_e2e tool to send the message 'hello'", "model": "{MODEL}", "provider": "{PROVIDER}"}
18+
"""
19+
Then The status code of the response is 200
20+
And The body of the response contains mock_tool_e2e

tests/e2e/secrets/mcp-token

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
test-secret-token

tests/e2e/test_list.txt

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,4 +15,5 @@ features/rlsapi_v1_errors.feature
1515
features/streaming_query.feature
1616
features/rest_api.feature
1717
features/mcp.feature
18+
features/mcp_file_auth.feature
1819
features/models.feature

tests/unit/utils/test_types.py

Lines changed: 132 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,13 +2,18 @@
22

33
import pytest
44
from llama_stack_api import ImageContentItem, TextContentItem, URL, _URLOrData
5+
from llama_stack_api.openai_responses import (
6+
OpenAIResponseInputToolFileSearch as InputToolFileSearch,
7+
OpenAIResponseInputToolMCP as InputToolMCP,
8+
)
59

610
from pydantic import AnyUrl, ValidationError
711
from pytest_mock import MockerFixture
812

913
from utils.types import (
1014
GraniteToolParser,
1115
ReferencedDocument,
16+
ResponsesApiParams,
1217
ToolCallSummary,
1318
ToolResultSummary,
1419
content_to_str,
@@ -194,3 +199,130 @@ def test_constructor_partial_fields(self) -> None:
194199
doc = ReferencedDocument(doc_title="Test Title")
195200
assert doc.doc_url is None
196201
assert doc.doc_title == "Test Title"
202+
203+
204+
class TestResponsesApiParamsModelDump:
205+
"""Tests for ResponsesApiParams.model_dump() MCP authorization fix.
206+
207+
Regression tests for LCORE-1414 / GitHub issue #1269: llama-stack-api's
208+
InputToolMCP.authorization has Field(exclude=True), causing the base
209+
model_dump() to silently strip authorization tokens.
210+
"""
211+
212+
def _make_params(self, tools: list) -> ResponsesApiParams:
213+
"""Build minimal ResponsesApiParams with given tools."""
214+
return ResponsesApiParams(
215+
input="test question",
216+
model="provider/model",
217+
conversation="conv-id",
218+
store=False,
219+
stream=False,
220+
tools=tools,
221+
)
222+
223+
def test_mcp_authorization_survives_model_dump(self) -> None:
224+
"""Test that MCP authorization is re-injected after model_dump()."""
225+
tool = InputToolMCP(
226+
server_label="auth-server",
227+
server_url="http://localhost:3000",
228+
require_approval="never",
229+
authorization="my-secret-token",
230+
)
231+
assert tool.authorization == "my-secret-token"
232+
assert "authorization" not in tool.model_dump()
233+
234+
params = self._make_params([tool])
235+
dumped = params.model_dump(exclude_none=True)
236+
assert dumped["tools"][0]["authorization"] == "my-secret-token"
237+
238+
def test_mcp_authorization_none_not_injected(self) -> None:
239+
"""Test that None authorization is not added to the dump."""
240+
tool = InputToolMCP(
241+
server_label="no-auth-server",
242+
server_url="http://localhost:3000",
243+
require_approval="never",
244+
)
245+
params = self._make_params([tool])
246+
dumped = params.model_dump(exclude_none=True)
247+
assert "authorization" not in dumped["tools"][0]
248+
249+
def test_non_mcp_tools_unaffected(self) -> None:
250+
"""Test that non-MCP tools are not modified by the override."""
251+
tool = InputToolFileSearch(
252+
type="file_search",
253+
vector_store_ids=["vs-1"],
254+
)
255+
params = self._make_params([tool])
256+
dumped = params.model_dump(exclude_none=True)
257+
assert "authorization" not in dumped["tools"][0]
258+
259+
def test_mixed_tools_only_mcp_gets_authorization(self) -> None:
260+
"""Test mixed tool list: only MCP tools get authorization re-injected."""
261+
mcp_tool = InputToolMCP(
262+
server_label="auth-server",
263+
server_url="http://localhost:3000",
264+
require_approval="never",
265+
authorization="secret",
266+
)
267+
file_tool = InputToolFileSearch(
268+
type="file_search",
269+
vector_store_ids=["vs-1"],
270+
)
271+
params = self._make_params([file_tool, mcp_tool])
272+
dumped = params.model_dump(exclude_none=True)
273+
274+
assert "authorization" not in dumped["tools"][0]
275+
assert dumped["tools"][1]["authorization"] == "secret"
276+
277+
def test_multiple_mcp_tools_each_preserves_authorization(self) -> None:
278+
"""Test that each MCP tool gets its own authorization re-injected."""
279+
tool_a = InputToolMCP(
280+
server_label="server-a",
281+
server_url="http://a:3000",
282+
require_approval="never",
283+
authorization="token-a",
284+
)
285+
tool_b = InputToolMCP(
286+
server_label="server-b",
287+
server_url="http://b:3000",
288+
require_approval="never",
289+
authorization="token-b",
290+
)
291+
params = self._make_params([tool_a, tool_b])
292+
dumped = params.model_dump(exclude_none=True)
293+
294+
assert dumped["tools"][0]["authorization"] == "token-a"
295+
assert dumped["tools"][1]["authorization"] == "token-b"
296+
297+
def test_exclude_changing_tool_list_shape_skips_reinjection(self) -> None:
298+
"""Test that exclude removing tool indices does not mis-assign authorization."""
299+
tool_a = InputToolMCP(
300+
server_label="server-a",
301+
server_url="http://a:3000",
302+
require_approval="never",
303+
authorization="token-a",
304+
)
305+
tool_b = InputToolMCP(
306+
server_label="server-b",
307+
server_url="http://b:3000",
308+
require_approval="never",
309+
authorization="token-b",
310+
)
311+
params = self._make_params([tool_a, tool_b])
312+
dumped = params.model_dump(exclude={"tools": {0}})
313+
assert len(dumped["tools"]) == 1
314+
assert dumped["tools"][0]["server_label"] == "server-b"
315+
assert "authorization" not in dumped["tools"][0]
316+
317+
def test_no_tools_does_not_error(self) -> None:
318+
"""Test that model_dump() works when tools is None."""
319+
params = ResponsesApiParams(
320+
input="test",
321+
model="provider/model",
322+
conversation="conv-id",
323+
store=False,
324+
stream=False,
325+
tools=None,
326+
)
327+
dumped = params.model_dump(exclude_none=True)
328+
assert "tools" not in dumped

0 commit comments

Comments
 (0)