From c6fa47db6f61c0e6da0ec8ce74b06fefa18e7225 Mon Sep 17 00:00:00 2001 From: Matthew Farrellee Date: Fri, 11 Apr 2025 12:25:57 -0400 Subject: [PATCH] fix: ensure resource registration arguments are typed (#1941) # What does this PR do? closes https://github.com/meta-llama/llama-stack/issues/1586 this issue arises when loading an mcp_endpoint from run.yaml. the issue does not manifest for mcp servers added via a running distro server. the existing tests only cover the case of adding to a running server. the code for loading run.yaml strips type information from mcp_endpoint, passing `{"uri": ...}` instead of `URL(uri=...)` along to the resource provider registration. ## Test Plan 1. run an mcp server 2. add an mcp tool config to the dev.py, e.g. ``` diff --git a/llama_stack/templates/dev/dev.py b/llama_stack/templates/dev/dev.py index 69924acb..e0dc7189 100644 --- a/llama_stack/templates/dev/dev.py +++ b/llama_stack/templates/dev/dev.py @@ -6,6 +6,8 @@ from typing import List, Tuple +from llama_stack.apis.common.content_types import URL + from llama_stack.apis.models.models import ModelType from llama_stack.distribution.datatypes import ( ModelInput, @@ -154,6 +156,11 @@ def get_distribution_template() -> DistributionTemplate: toolgroup_id="builtin::code_interpreter", provider_id="code-interpreter", ), + ToolGroupInput( + toolgroup_id="mcp::filesystem", + provider_id="model-context-protocol", + mcp_endpoint=URL(uri="http://localhost:8002/sse"), + ), ] embedding_model = ModelInput( model_id="all-MiniLM-L6-v2", ``` 3. run distro_codegen.py 4. llama stack build --template dev --run before this pr, the `llama stack run` would fail w/ `AttributeError: 'dict' object has no attribute 'uri'`, after it will succeed. --- llama_stack/distribution/stack.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/llama_stack/distribution/stack.py b/llama_stack/distribution/stack.py index d70878db4..08ff5e7cd 100644 --- a/llama_stack/distribution/stack.py +++ b/llama_stack/distribution/stack.py @@ -96,7 +96,10 @@ async def register_resources(run_config: StackRunConfig, impls: Dict[Api, Any]): method = getattr(impls[api], register_method) for obj in objects: - await method(**obj.model_dump()) + # we want to maintain the type information in arguments to method. + # instead of method(**obj.model_dump()), which may convert a typed attr to a dict, + # we use model_dump() to find all the attrs and then getattr to get the still typed value. + await method(**{k: getattr(obj, k) for k in obj.model_dump().keys()}) method = getattr(impls[api], list_method) response = await method()