llama-stack 0.4.2__py3-none-any.whl → 0.4.3__py3-none-any.whl

llama_stack/core/library_client.py

@@ -161,6 +161,45 @@ class LlamaStackAsLibraryClient(LlamaStackClient):
         """
         pass
 
+    def shutdown(self) -> None:
+        """Shutdown the client and release all resources.
+
+        This method should be called when you're done using the client to properly
+        close database connections and release other resources. Failure to call this
+        method may result in the program hanging on exit while waiting for background
+        threads to complete.
+
+        This method is idempotent and can be called multiple times safely.
+
+        Example:
+            client = LlamaStackAsLibraryClient("starter")
+            # ... use the client ...
+            client.shutdown()
+        """
+        loop = self.loop
+        asyncio.set_event_loop(loop)
+        try:
+            loop.run_until_complete(self.async_client.shutdown())
+        finally:
+            loop.close()
+            asyncio.set_event_loop(None)
+
+    def __enter__(self) -> "LlamaStackAsLibraryClient":
+        """Enter the context manager.
+
+        The client is already initialized in __init__, so this just returns self.
+
+        Example:
+            with LlamaStackAsLibraryClient("starter") as client:
+                response = client.models.list()
+            # Client is automatically shut down here
+        """
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
+        """Exit the context manager and shut down the client."""
+        self.shutdown()
+
     def request(self, *args, **kwargs):
         loop = self.loop
         asyncio.set_event_loop(loop)
@@ -224,6 +263,7 @@ class AsyncLlamaStackAsLibraryClient(AsyncLlamaStackClient):
         self.custom_provider_registry = custom_provider_registry
         self.provider_data = provider_data
         self.route_impls: RouteImpls | None = None  # Initialize to None to prevent AttributeError
+        self.stack: Stack | None = None
 
     def _remove_root_logger_handlers(self):
         """
@@ -246,9 +286,9 @@ class AsyncLlamaStackAsLibraryClient(AsyncLlamaStackClient):
         try:
             self.route_impls = None
 
-            stack = Stack(self.config, self.custom_provider_registry)
-            await stack.initialize()
-            self.impls = stack.impls
+            self.stack = Stack(self.config, self.custom_provider_registry)
+            await self.stack.initialize()
+            self.impls = self.stack.impls
         except ModuleNotFoundError as _e:
             cprint(_e.msg, color="red", file=sys.stderr)
             cprint(
@@ -283,6 +323,43 @@ class AsyncLlamaStackAsLibraryClient(AsyncLlamaStackClient):
         self.route_impls = initialize_route_impls(self.impls)
         return True
 
+    async def shutdown(self) -> None:
+        """Shutdown the client and release all resources.
+
+        This method should be called when you're done using the client to properly
+        close database connections and release other resources. Failure to call this
+        method may result in the program hanging on exit while waiting for background
+        threads to complete.
+
+        This method is idempotent and can be called multiple times safely.
+
+        Example:
+            client = AsyncLlamaStackAsLibraryClient("starter")
+            await client.initialize()
+            # ... use the client ...
+            await client.shutdown()
+        """
+        if self.stack:
+            await self.stack.shutdown()
+            self.stack = None
+
+    async def __aenter__(self) -> "AsyncLlamaStackAsLibraryClient":
+        """Enter the async context manager.
+
+        Initializes the client and returns it.
+
+        Example:
+            async with AsyncLlamaStackAsLibraryClient("starter") as client:
+                response = await client.models.list()
+            # Client is automatically shut down here
+        """
+        await self.initialize()
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
+        """Exit the async context manager and shut down the client."""
+        await self.shutdown()
+
     async def request(
         self,
         cast_to: Any,

llama_stack/core/routing_tables/common.py

@@ -209,6 +209,17 @@ class CommonRoutingTableImpl(RoutingTable):
             logger.info(f"Setting owner for {obj.type} '{obj.identifier}' to {obj.owner.principal}")
 
         registered_obj = await register_object_with_provider(obj, p)
+
+        # Ensure OpenAI metadata exists for vector stores
+        if obj.type == ResourceType.vector_store.value:
+            if hasattr(p, "_ensure_openai_metadata_exists"):
+                await p._ensure_openai_metadata_exists(obj)
+            else:
+                logger.warning(
+                    f"Provider {obj.provider_id} does not support OpenAI metadata creation. "
+                    f"Vector store {obj.identifier} may not work with OpenAI-compatible APIs."
+                )
+
         # TODO: This needs to be fixed for all APIs once they return the registered object
         if obj.type == ResourceType.model.value:
             await self.dist_registry.register(registered_obj)

llama_stack/core/routing_tables/vector_stores.py

@@ -55,6 +55,10 @@ class VectorStoresRoutingTable(CommonRoutingTableImpl):
 
     # Internal methods only - no public API exposure
 
+    async def list_vector_stores(self) -> list[VectorStoreWithOwner]:
+        """List all registered vector stores."""
+        return await self.get_all_with_type(ResourceType.vector_store.value)
+
     async def register_vector_store(
         self,
         vector_store_id: str,

llama_stack/core/stack.py

@@ -108,6 +108,7 @@ RESOURCES = [
     ),
     ("benchmarks", Api.benchmarks, "register_benchmark", "list_benchmarks", RegisterBenchmarkRequest),
     ("tool_groups", Api.tool_groups, "register_tool_group", "list_tool_groups", None),
+    ("vector_stores", Api.vector_stores, "register_vector_store", "list_vector_stores", None),
 ]
 
 
@@ -620,7 +621,7 @@ class Stack:
     async def shutdown(self):
         for impl in self.impls.values():
             impl_name = impl.__class__.__name__
-            logger.info(f"Shutting down {impl_name}")
+            logger.debug(f"Shutting down {impl_name}")
             try:
                 if hasattr(impl, "shutdown"):
                     await asyncio.wait_for(impl.shutdown(), timeout=5)
@@ -642,6 +643,20 @@ class Stack:
         if REGISTRY_REFRESH_TASK:
             REGISTRY_REFRESH_TASK.cancel()
 
+        # Shutdown storage backends
+        from llama_stack.core.storage.kvstore.kvstore import shutdown_kvstore_backends
+        from llama_stack.core.storage.sqlstore.sqlstore import shutdown_sqlstore_backends
+
+        try:
+            await shutdown_kvstore_backends()
+        except Exception as e:
+            logger.exception(f"Failed to shutdown KV store backends: {e}")
+
+        try:
+            await shutdown_sqlstore_backends()
+        except Exception as e:
+            logger.exception(f"Failed to shutdown SQL store backends: {e}")
+
 
 async def refresh_registry_once(impls: dict[Api, Any]):
     logger.debug("refreshing registry")

llama_stack/core/storage/kvstore/kvstore.py

@@ -62,6 +62,9 @@ class InmemoryKVStoreImpl(KVStore):
     async def delete(self, key: str) -> None:
         del self._store[key]
 
+    async def shutdown(self) -> None:
+        self._store.clear()
+
 
 _KVSTORE_BACKENDS: dict[str, KVStoreConfig] = {}
 _KVSTORE_INSTANCES: dict[tuple[str, str], KVStore] = {}
@@ -126,3 +129,11 @@ async def kvstore_impl(reference: KVStoreReference) -> KVStore:
         await impl.initialize()
         _KVSTORE_INSTANCES[cache_key] = impl
         return impl
+
+
+async def shutdown_kvstore_backends() -> None:
+    """Shutdown all cached KV store instances."""
+    global _KVSTORE_INSTANCES
+    for instance in _KVSTORE_INSTANCES.values():
+        await instance.shutdown()
+    _KVSTORE_INSTANCES.clear()

llama_stack/core/storage/kvstore/mongodb/mongodb.py

@@ -83,3 +83,8 @@ class MongoDBKVStoreImpl(KVStore):
         async for doc in cursor:
             result.append(doc["key"])
         return result
+
+    async def shutdown(self) -> None:
+        if self.conn:
+            await self.conn.close()
+            self.conn = None

llama_stack/core/storage/kvstore/postgres/postgres.py

@@ -123,3 +123,11 @@ class PostgresKVStoreImpl(KVStore):
             (start_key, end_key),
         )
         return [row[0] for row in cursor.fetchall()]
+
+    async def shutdown(self) -> None:
+        if self._cursor:
+            self._cursor.close()
+            self._cursor = None
+        if self._conn:
+            self._conn.close()
+            self._conn = None

llama_stack/core/storage/kvstore/redis/redis.py

@@ -99,3 +99,8 @@ class RedisKVStoreImpl(KVStore):
             if cursor == 0:
                 break
         return result
+
+    async def shutdown(self) -> None:
+        if self._redis:
+            await self._redis.close()
+            self._redis = None

llama_stack/core/storage/sqlstore/sqlalchemy_sqlstore.py

@@ -107,6 +107,14 @@ class SqlAlchemySqlStoreImpl(SqlStore):
 
         return engine
 
+    async def shutdown(self) -> None:
+        """Dispose the session maker's engine and close all connections."""
+        # The async_session holds a reference to the engine created in __init__
+        if self.async_session:
+            engine = self.async_session.kw.get("bind")
+            if engine:
+                await engine.dispose()
+
     async def create_table(
         self,
         table: str,

llama_stack/core/storage/sqlstore/sqlstore.py

@@ -85,3 +85,11 @@ def register_sqlstore_backends(backends: dict[str, StorageBackendConfig]) -> Non
     _SQLSTORE_LOCKS.clear()
     for name, cfg in backends.items():
         _SQLSTORE_BACKENDS[name] = cfg
+
+
+async def shutdown_sqlstore_backends() -> None:
+    """Shutdown all cached SQL store instances."""
+    global _SQLSTORE_INSTANCES
+    for instance in _SQLSTORE_INSTANCES.values():
+        await instance.shutdown()
+    _SQLSTORE_INSTANCES.clear()

llama_stack/providers/inline/agents/meta_reference/responses/openai_responses.py

@@ -4,6 +4,7 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 
+import asyncio
 import re
 import time
 import uuid
@@ -16,6 +17,7 @@ from llama_stack.providers.utils.responses.responses_store import (
     ResponsesStore,
     _OpenAIResponseObjectWithInputAndMessages,
 )
+from llama_stack.providers.utils.tools.mcp import MCPSessionManager
 from llama_stack_api import (
     ConversationItem,
     Conversations,
@@ -489,6 +491,19 @@ class OpenAIResponsesImpl:
         response_id = f"resp_{uuid.uuid4()}"
         created_at = int(time.time())
 
+        # Create a per-request MCP session manager for session reuse (fix for #4452)
+        # This avoids redundant tools/list calls when making multiple MCP tool invocations
+        mcp_session_manager = MCPSessionManager()
+
+        # Create a per-request ToolExecutor with the session manager
+        request_tool_executor = ToolExecutor(
+            tool_groups_api=self.tool_groups_api,
+            tool_runtime_api=self.tool_runtime_api,
+            vector_io_api=self.vector_io_api,
+            vector_stores_config=self.tool_executor.vector_stores_config,
+            mcp_session_manager=mcp_session_manager,
+        )
+
         orchestrator = StreamingResponseOrchestrator(
             inference_api=self.inference_api,
             ctx=ctx,
@@ -498,7 +513,7 @@ class OpenAIResponsesImpl:
             text=text,
             max_infer_iters=max_infer_iters,
             parallel_tool_calls=parallel_tool_calls,
-            tool_executor=self.tool_executor,
+            tool_executor=request_tool_executor,
             safety_api=self.safety_api,
             guardrail_ids=guardrail_ids,
             instructions=instructions,
@@ -513,41 +528,52 @@ class OpenAIResponsesImpl:
 
         # Type as ConversationItem to avoid list invariance issues
         output_items: list[ConversationItem] = []
-        async for stream_chunk in orchestrator.create_response():
-            match stream_chunk.type:
-                case "response.completed" | "response.incomplete":
-                    final_response = stream_chunk.response
-                case "response.failed":
-                    failed_response = stream_chunk.response
-                case "response.output_item.done":
-                    item = stream_chunk.item
-                    output_items.append(item)
-                case _:
-                    pass  # Other event types
-
-            # Store and sync before yielding terminal events
-            # This ensures the storage/syncing happens even if the consumer breaks after receiving the event
-            if (
-                stream_chunk.type in {"response.completed", "response.incomplete"}
-                and final_response
-                and failed_response is None
-            ):
-                messages_to_store = list(
-                    filter(lambda x: not isinstance(x, OpenAISystemMessageParam), orchestrator.final_messages)
-                )
-                if store:
-                    # TODO: we really should work off of output_items instead of "final_messages"
-                    await self._store_response(
-                        response=final_response,
-                        input=all_input,
-                        messages=messages_to_store,
+        try:
+            async for stream_chunk in orchestrator.create_response():
+                match stream_chunk.type:
+                    case "response.completed" | "response.incomplete":
+                        final_response = stream_chunk.response
+                    case "response.failed":
+                        failed_response = stream_chunk.response
+                    case "response.output_item.done":
+                        item = stream_chunk.item
+                        output_items.append(item)
+                    case _:
+                        pass  # Other event types
+
+                # Store and sync before yielding terminal events
+                # This ensures the storage/syncing happens even if the consumer breaks after receiving the event
+                if (
+                    stream_chunk.type in {"response.completed", "response.incomplete"}
+                    and final_response
+                    and failed_response is None
+                ):
+                    messages_to_store = list(
+                        filter(lambda x: not isinstance(x, OpenAISystemMessageParam), orchestrator.final_messages)
                     )
+                    if store:
+                        # TODO: we really should work off of output_items instead of "final_messages"
+                        await self._store_response(
+                            response=final_response,
+                            input=all_input,
+                            messages=messages_to_store,
+                        )
 
-                if conversation:
-                    await self._sync_response_to_conversation(conversation, input, output_items)
-                    await self.responses_store.store_conversation_messages(conversation, messages_to_store)
-
-            yield stream_chunk
+                    if conversation:
+                        await self._sync_response_to_conversation(conversation, input, output_items)
+                        await self.responses_store.store_conversation_messages(conversation, messages_to_store)
+
+                yield stream_chunk
+        finally:
+            # Clean up MCP sessions at the end of the request (fix for #4452)
+            # Use shield() to prevent cancellation from interrupting cleanup and leaking resources
+            # Wrap in try/except as cleanup errors should not mask the original response
+            try:
+                await asyncio.shield(mcp_session_manager.close_all())
+            except BaseException as e:
+                # Debug level - cleanup errors are expected in streaming scenarios where
+                # anyio cancel scopes may be in a different task context
+                logger.debug(f"Error during MCP session cleanup: {e}")
 
     async def delete_openai_response(self, response_id: str) -> OpenAIDeleteResponseObject:
         return await self.responses_store.delete_response_object(response_id)

llama_stack/providers/inline/agents/meta_reference/responses/streaming.py

@@ -1200,6 +1200,9 @@ class StreamingResponseOrchestrator:
                 "mcp_list_tools_id": list_id,
             }
 
+            # Get session manager from tool_executor if available (fix for #4452)
+            session_manager = getattr(self.tool_executor, "mcp_session_manager", None)
+
             # TODO: follow semantic conventions for Open Telemetry tool spans
             # https://opentelemetry.io/docs/specs/semconv/gen-ai/gen-ai-spans/#execute-tool-span
             with tracer.start_as_current_span("list_mcp_tools", attributes=attributes):
@@ -1207,6 +1210,7 @@ class StreamingResponseOrchestrator:
                     endpoint=mcp_tool.server_url,
                     headers=mcp_tool.headers,
                     authorization=mcp_tool.authorization,
+                    session_manager=session_manager,
                 )
 
             # Create the MCP list tools message

llama_stack/providers/inline/agents/meta_reference/responses/tool_executor.py

@@ -54,11 +54,14 @@ class ToolExecutor:
         tool_runtime_api: ToolRuntime,
         vector_io_api: VectorIO,
         vector_stores_config=None,
+        mcp_session_manager=None,
     ):
         self.tool_groups_api = tool_groups_api
         self.tool_runtime_api = tool_runtime_api
         self.vector_io_api = vector_io_api
         self.vector_stores_config = vector_stores_config
+        # Optional MCPSessionManager for session reuse within a request (fix for #4452)
+        self.mcp_session_manager = mcp_session_manager
 
     async def execute_tool_call(
         self,
@@ -233,6 +236,7 @@ class ToolExecutor:
                 "document_ids": [r.file_id for r in search_results],
                 "chunks": [r.content[0].text if r.content else "" for r in search_results],
                 "scores": [r.score for r in search_results],
+                "attributes": [r.attributes or {} for r in search_results],
                 "citation_files": citation_files,
             },
         )
@@ -327,12 +331,14 @@ class ToolExecutor:
                 # TODO: follow semantic conventions for Open Telemetry tool spans
                 # https://opentelemetry.io/docs/specs/semconv/gen-ai/gen-ai-spans/#execute-tool-span
                 with tracer.start_as_current_span("invoke_mcp_tool", attributes=attributes):
+                    # Pass session_manager for session reuse within request (fix for #4452)
                     result = await invoke_mcp_tool(
                         endpoint=mcp_tool.server_url,
                         tool_name=function_name,
                         kwargs=tool_kwargs,
                         headers=mcp_tool.headers,
                         authorization=mcp_tool.authorization,
+                        session_manager=self.mcp_session_manager,
                     )
             elif function_name == "knowledge_search":
                 response_file_search_tool = (
@@ -464,16 +470,18 @@ class ToolExecutor:
                 )
                 if result and (metadata := getattr(result, "metadata", None)) and "document_ids" in metadata:
                     message.results = []
+                    attributes_list = metadata.get("attributes", [])
                     for i, doc_id in enumerate(metadata["document_ids"]):
                         text = metadata["chunks"][i] if "chunks" in metadata else None
                         score = metadata["scores"][i] if "scores" in metadata else None
+                        attrs = attributes_list[i] if i < len(attributes_list) else {}
                         message.results.append(
                             OpenAIResponseOutputMessageFileSearchToolCallResults(
                                 file_id=doc_id,
                                 filename=doc_id,
                                 text=text if text is not None else "",
                                 score=score if score is not None else 0.0,
-                                attributes={},
+                                attributes=attrs,
                             )
                         )
                 if has_error:

llama_stack/providers/remote/vector_io/pgvector/pgvector.py

@@ -10,6 +10,7 @@ from typing import Any
 import psycopg2
 from numpy.typing import NDArray
 from psycopg2 import sql
+from psycopg2.extensions import cursor
 from psycopg2.extras import Json, execute_values
 from pydantic import BaseModel, TypeAdapter
 
@@ -54,6 +55,17 @@ def check_extension_version(cur):
     return result[0] if result else None
 
 
+def create_vector_extension(cur: cursor) -> None:
+    try:
+        log.info("Vector extension not found, creating...")
+        cur.execute("CREATE EXTENSION vector;")
+        log.info("Vector extension created successfully")
+        log.info(f"Vector extension version: {check_extension_version(cur)}")
+
+    except psycopg2.Error as e:
+        raise RuntimeError(f"Failed to create vector extension for PGVector: {e}") from e
+
+
 def upsert_models(conn, keys_models: list[tuple[str, BaseModel]]):
     with conn.cursor(cursor_factory=psycopg2.extras.DictCursor) as cur:
         query = sql.SQL(
@@ -364,7 +376,7 @@ class PGVectorVectorIOAdapter(OpenAIVectorStoreMixin, VectorIO, VectorStoresProt
                 if version:
                     log.info(f"Vector extension version: {version}")
                 else:
-                    raise RuntimeError("Vector extension is not installed.")
+                    create_vector_extension(cur)
 
                 cur.execute(
                     """

llama_stack/providers/utils/inference/embedding_mixin.py

@@ -25,7 +25,8 @@ from llama_stack_api import (
     OpenAIEmbeddingUsage,
 )
 
-EMBEDDING_MODELS = {}
+EMBEDDING_MODELS: dict[str, "SentenceTransformer"] = {}
+EMBEDDING_MODELS_LOCK = asyncio.Lock()
 
 DARWIN = "Darwin"
 
@@ -76,26 +77,29 @@ class SentenceTransformerEmbeddingMixin:
         )
 
     async def _load_sentence_transformer_model(self, model: str) -> "SentenceTransformer":
-        global EMBEDDING_MODELS
-
         loaded_model = EMBEDDING_MODELS.get(model)
         if loaded_model is not None:
             return loaded_model
 
-        log.info(f"Loading sentence transformer for {model}...")
+        async with EMBEDDING_MODELS_LOCK:
+            loaded_model = EMBEDDING_MODELS.get(model)
+            if loaded_model is not None:
+                return loaded_model
+
+            log.info(f"Loading sentence transformer for {model}...")
 
-        def _load_model():
-            from sentence_transformers import SentenceTransformer
+            def _load_model():
+                from sentence_transformers import SentenceTransformer
 
-            platform_name = platform.system()
-            if platform_name == DARWIN:
-                # PyTorch's OpenMP kernels can segfault on macOS when spawned from background
-                # threads with the default parallel settings, so force a single-threaded CPU run.
-                log.debug(f"Constraining torch threads on {platform_name} to a single worker")
-                torch.set_num_threads(1)
+                platform_name = platform.system()
+                if platform_name == DARWIN:
+                    # PyTorch's OpenMP kernels can segfault on macOS when spawned from background
+                    # threads with the default parallel settings, so force a single-threaded CPU run.
+                    log.debug(f"Constraining torch threads on {platform_name} to a single worker")
+                    torch.set_num_threads(1)
 
-            return SentenceTransformer(model, trust_remote_code=True)
+                return SentenceTransformer(model, trust_remote_code=True)
 
-        loaded_model = await asyncio.to_thread(_load_model)
-        EMBEDDING_MODELS[model] = loaded_model
-        return loaded_model
+            loaded_model = await asyncio.to_thread(_load_model)
+            EMBEDDING_MODELS[model] = loaded_model
+            return loaded_model

llama_stack/providers/utils/memory/openai_vector_store_mixin.py

@@ -122,6 +122,39 @@ class OpenAIVectorStoreMixin(ABC):
         # update in-memory cache
         self.openai_vector_stores[store_id] = store_info
 
+    async def _ensure_openai_metadata_exists(self, vector_store: VectorStore, name: str | None = None) -> None:
+        """
+        Ensure OpenAI-compatible metadata exists for a vector store.
+        """
+        if vector_store.identifier not in self.openai_vector_stores:
+            store_info = {
+                "id": vector_store.identifier,
+                "object": "vector_store",
+                "created_at": int(time.time()),
+                "name": name or vector_store.vector_store_name or vector_store.identifier,
+                "usage_bytes": 0,
+                "file_counts": VectorStoreFileCounts(
+                    cancelled=0,
+                    completed=0,
+                    failed=0,
+                    in_progress=0,
+                    total=0,
+                ).model_dump(),
+                "status": "completed",
+                "expires_after": None,
+                "expires_at": None,
+                "last_active_at": int(time.time()),
+                "file_ids": [],
+                "chunking_strategy": None,
+                "metadata": {
+                    "provider_id": vector_store.provider_id,
+                    "provider_vector_store_id": vector_store.provider_resource_id,
+                    "embedding_model": vector_store.embedding_model,
+                    "embedding_dimension": str(vector_store.embedding_dimension),
+                },
+            }
+            await self._save_openai_vector_store(vector_store.identifier, store_info)
+
     async def _load_openai_vector_stores(self) -> dict[str, dict[str, Any]]:
         """Load all vector store metadata from persistent storage."""
         assert self.kvstore

llama_stack/providers/utils/tools/mcp.py

@@ -4,6 +4,8 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 
+import asyncio
+import hashlib
 from collections.abc import AsyncGenerator
 from contextlib import asynccontextmanager
 from enum import Enum
@@ -73,6 +75,207 @@ class MCPProtol(Enum):
     SSE = 2
 
 
+class MCPSessionManager:
+    """Manages MCP session lifecycle within a request scope.
+
+    This class caches MCP sessions by (endpoint, headers_hash) to avoid redundant
+    connection establishment and tools/list calls when making multiple tool
+    invocations to the same MCP server within a single request.
+
+    Fix for GitHub issue #4452: MCP tools/list called redundantly before every
+    tool invocation.
+
+    Usage:
+        async with MCPSessionManager() as session_manager:
+            # Multiple tool calls will reuse the same session
+            result1 = await invoke_mcp_tool(..., session_manager=session_manager)
+            result2 = await invoke_mcp_tool(..., session_manager=session_manager)
+    """
+
+    def __init__(self):
+        # Cache of active sessions: key -> (session, client_context, session_context)
+        self._sessions: dict[str, tuple[ClientSession, Any, Any]] = {}
+        # Locks to prevent concurrent session creation for the same key
+        self._locks: dict[str, asyncio.Lock] = {}
+        # Global lock for managing the locks dict
+        self._global_lock = asyncio.Lock()
+
+    def _make_key(self, endpoint: str, headers: dict[str, str]) -> str:
+        """Create a cache key from endpoint and headers."""
+        # Sort headers for consistent hashing
+        headers_str = str(sorted(headers.items()))
+        headers_hash = hashlib.sha256(headers_str.encode()).hexdigest()[:16]
+        return f"{endpoint}:{headers_hash}"
+
+    async def _get_lock(self, key: str) -> asyncio.Lock:
+        """Get or create a lock for a specific cache key."""
+        async with self._global_lock:
+            if key not in self._locks:
+                self._locks[key] = asyncio.Lock()
+            return self._locks[key]
+
+    async def get_session(self, endpoint: str, headers: dict[str, str]) -> ClientSession:
+        """Get or create an MCP session for the given endpoint and headers.
+
+        Args:
+            endpoint: MCP server endpoint URL
+            headers: Headers including authorization
+
+        Returns:
+            An initialized ClientSession ready for tool calls
+        """
+        key = self._make_key(endpoint, headers)
+
+        # Check if session already exists (fast path)
+        if key in self._sessions:
+            session, _, _ = self._sessions[key]
+            return session
+
+        # Acquire lock for this specific key to prevent concurrent creation
+        lock = await self._get_lock(key)
+        async with lock:
+            # Double-check after acquiring lock
+            if key in self._sessions:
+                session, _, _ = self._sessions[key]
+                return session
+
+            # Create new session
+            session, client_ctx, session_ctx = await self._create_session(endpoint, headers)
+            self._sessions[key] = (session, client_ctx, session_ctx)
+            logger.debug(f"Created new MCP session for {endpoint} (key: {key[:32]}...)")
+            return session
+
+    async def _create_session(self, endpoint: str, headers: dict[str, str]) -> tuple[ClientSession, Any, Any]:
+        """Create a new MCP session.
+
+        Returns:
+            Tuple of (session, client_context, session_context) for lifecycle management
+        """
+        # Use the same protocol detection logic as client_wrapper
+        connection_strategies = [MCPProtol.STREAMABLE_HTTP, MCPProtol.SSE]
+        mcp_protocol = protocol_cache.get(endpoint, default=MCPProtol.UNKNOWN)
+        if mcp_protocol == MCPProtol.SSE:
+            connection_strategies = [MCPProtol.SSE, MCPProtol.STREAMABLE_HTTP]
+
+        last_exception: BaseException | None = None
+
+        for i, strategy in enumerate(connection_strategies):
+            try:
+                client = streamablehttp_client
+                if strategy == MCPProtol.SSE:
+                    client = cast(Any, sse_client)
+
+                # Enter the client context manager manually
+                client_ctx = client(endpoint, headers=headers)
+                client_streams = await client_ctx.__aenter__()
+
+                try:
+                    # Enter the session context manager manually
+                    session = ClientSession(read_stream=client_streams[0], write_stream=client_streams[1])
+                    session_ctx = session
+                    await session.__aenter__()
+
+                    try:
+                        await session.initialize()
+                        protocol_cache[endpoint] = strategy
+                        return session, client_ctx, session_ctx
+                    except BaseException:
+                        await session.__aexit__(None, None, None)
+                        raise
+                except BaseException:
+                    await client_ctx.__aexit__(None, None, None)
+                    raise
+
+            except* httpx.HTTPStatusError as eg:
+                for exc in eg.exceptions:
+                    err = cast(httpx.HTTPStatusError, exc)
+                    if err.response.status_code == 401:
+                        raise AuthenticationRequiredError(exc) from exc
+                if i == len(connection_strategies) - 1:
+                    raise
+                last_exception = eg
+            except* httpx.ConnectError as eg:
+                if i == len(connection_strategies) - 1:
+                    error_msg = f"Failed to connect to MCP server at {endpoint}: Connection refused"
+                    logger.error(f"MCP connection error: {error_msg}")
+                    raise ConnectionError(error_msg) from eg
+                else:
+                    logger.warning(
+                        f"failed to connect to MCP server at {endpoint} via {strategy.name}, "
+                        f"falling back to {connection_strategies[i + 1].name}"
+                    )
+                last_exception = eg
+            except* httpx.TimeoutException as eg:
+                if i == len(connection_strategies) - 1:
+                    error_msg = f"MCP server at {endpoint} timed out"
+                    logger.error(f"MCP timeout error: {error_msg}")
+                    raise TimeoutError(error_msg) from eg
+                else:
+                    logger.warning(
+                        f"MCP server at {endpoint} timed out via {strategy.name}, "
+                        f"falling back to {connection_strategies[i + 1].name}"
+                    )
+                last_exception = eg
+            except* httpx.RequestError as eg:
+                if i == len(connection_strategies) - 1:
+                    exc_msg = str(eg.exceptions[0]) if eg.exceptions else "Unknown error"
+                    error_msg = f"Network error connecting to MCP server at {endpoint}: {exc_msg}"
+                    logger.error(f"MCP network error: {error_msg}")
+                    raise ConnectionError(error_msg) from eg
+                else:
+                    logger.warning(
+                        f"network error connecting to MCP server at {endpoint} via {strategy.name}, "
+                        f"falling back to {connection_strategies[i + 1].name}"
+                    )
+                last_exception = eg
+            except* McpError:
+                if i < len(connection_strategies) - 1:
+                    logger.warning(
+                        f"failed to connect via {strategy.name}, falling back to {connection_strategies[i + 1].name}"
+                    )
+                else:
+                    raise
+
+        # Should not reach here, but just in case
+        if last_exception:
+            raise last_exception
+        raise RuntimeError(f"Failed to create MCP session for {endpoint}")
+
+    async def close_all(self) -> None:
+        """Close all cached sessions.
+
+        Should be called at the end of a request to clean up resources.
+
+        Note: We catch BaseException (not just Exception) because:
+        1. CancelledError is a BaseException and can occur during cleanup
+        2. anyio cancel scope errors can occur if cleanup runs in a different
+           task context than where the session was created
+        These are expected in streaming response scenarios and are handled gracefully.
+        """
+        errors = []
+        session_count = len(self._sessions)
+        for key, (session, client_ctx, _) in list(self._sessions.items()):
+            try:
+                await session.__aexit__(None, None, None)
+            except BaseException as e:
+                # Debug level since these errors are expected in streaming scenarios
+                # where cleanup runs in a different async context than session creation
+                logger.debug(f"Error closing MCP session {key}: {e}")
+                errors.append(e)
+            try:
+                await client_ctx.__aexit__(None, None, None)
+            except BaseException as e:
+                logger.debug(f"Error closing MCP client context {key}: {e}")
+                errors.append(e)
+
+        self._sessions.clear()
+        self._locks.clear()
+        logger.debug(f"Closed {session_count} MCP sessions")
+
+        if errors:
+            logger.debug(f"Encountered {len(errors)} errors while closing MCP sessions (expected in streaming)")
+
+
 @asynccontextmanager
 async def client_wrapper(endpoint: str, headers: dict[str, str]) -> AsyncGenerator[ClientSession, Any]:
     # we use a ttl'd dict to cache the happy path protocol for each endpoint
@@ -151,6 +354,7 @@ async def list_mcp_tools(
     endpoint: str,
     headers: dict[str, str] | None = None,
     authorization: str | None = None,
+    session_manager: MCPSessionManager | None = None,
 ) -> ListToolDefsResponse:
     """List tools available from an MCP server.
 
@@ -158,6 +362,10 @@ async def list_mcp_tools(
         endpoint: MCP server endpoint URL
         headers: Optional base headers to include
         authorization: Optional OAuth access token (just the token, not "Bearer <token>")
+        session_manager: Optional MCPSessionManager for session reuse within a request.
+            When provided, sessions are cached and reused, avoiding redundant session
+            creation when list_mcp_tools and invoke_mcp_tool are called for the same
+            server within a request. (Fix for #4452)
 
     Returns:
         List of tool definitions from the MCP server
@@ -169,7 +377,9 @@ async def list_mcp_tools(
     final_headers = prepare_mcp_headers(headers, authorization)
 
     tools = []
-    async with client_wrapper(endpoint, final_headers) as session:
+
+    # Helper function to process session and list tools
+    async def _list_tools_from_session(session):
         tools_result = await session.list_tools()
         for tool in tools_result.tools:
             tools.append(
@@ -183,15 +393,51 @@ async def list_mcp_tools(
                     },
                 )
             )
+
+    # If a session manager is provided, use it for session reuse (fix for #4452)
+    if session_manager is not None:
+        session = await session_manager.get_session(endpoint, final_headers)
+        await _list_tools_from_session(session)
+    else:
+        # Fallback to original behavior: create a new session for this call
+        async with client_wrapper(endpoint, final_headers) as session:
+            await _list_tools_from_session(session)
+
     return ListToolDefsResponse(data=tools)
 
 
+def _parse_mcp_result(result) -> ToolInvocationResult:
+    """Parse MCP tool call result into ToolInvocationResult.
+
+    Args:
+        result: The raw MCP tool call result
+
+    Returns:
+        ToolInvocationResult with parsed content
+    """
+    content: list[InterleavedContentItem] = []
+    for item in result.content:
+        if isinstance(item, mcp_types.TextContent):
+            content.append(TextContentItem(text=item.text))
+        elif isinstance(item, mcp_types.ImageContent):
+            content.append(ImageContentItem(image=_URLOrData(data=item.data)))
+        elif isinstance(item, mcp_types.EmbeddedResource):
+            logger.warning(f"EmbeddedResource is not supported: {item}")
+        else:
+            raise ValueError(f"Unknown content type: {type(item)}")
+    return ToolInvocationResult(
+        content=content,
+        error_code=1 if result.isError else 0,
+    )
+
+
 async def invoke_mcp_tool(
     endpoint: str,
     tool_name: str,
     kwargs: dict[str, Any],
     headers: dict[str, str] | None = None,
     authorization: str | None = None,
+    session_manager: MCPSessionManager | None = None,
 ) -> ToolInvocationResult:
     """Invoke an MCP tool with the given arguments.
 
@@ -201,6 +447,9 @@ async def invoke_mcp_tool(
         kwargs: Tool invocation arguments
         headers: Optional base headers to include
         authorization: Optional OAuth access token (just the token, not "Bearer <token>")
+        session_manager: Optional MCPSessionManager for session reuse within a request.
+            When provided, sessions are cached and reused for multiple tool calls to
+            the same endpoint, avoiding redundant tools/list calls. (Fix for #4452)
 
     Returns:
         Tool invocation result with content and error information
@@ -211,20 +460,13 @@ async def invoke_mcp_tool(
     # Prepare headers with authorization handling
     final_headers = prepare_mcp_headers(headers, authorization)
 
-    async with client_wrapper(endpoint, final_headers) as session:
+    # If a session manager is provided, use it for session reuse (fix for #4452)
+    if session_manager is not None:
+        session = await session_manager.get_session(endpoint, final_headers)
         result = await session.call_tool(tool_name, kwargs)
+        return _parse_mcp_result(result)
 
-        content: list[InterleavedContentItem] = []
-        for item in result.content:
-            if isinstance(item, mcp_types.TextContent):
-                content.append(TextContentItem(text=item.text))
-            elif isinstance(item, mcp_types.ImageContent):
-                content.append(ImageContentItem(image=_URLOrData(data=item.data)))
-            elif isinstance(item, mcp_types.EmbeddedResource):
-                logger.warning(f"EmbeddedResource is not supported: {item}")
-            else:
-                raise ValueError(f"Unknown content type: {type(item)}")
-        return ToolInvocationResult(
-            content=content,
-            error_code=1 if result.isError else 0,
-        )
+    # Fallback to original behavior: create a new session for each call
+    async with client_wrapper(endpoint, final_headers) as session:
+        result = await session.call_tool(tool_name, kwargs)
+        return _parse_mcp_result(result)

{llama_stack-0.4.2.dist-info → llama_stack-0.4.3.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: llama_stack
-Version: 0.4.2
+Version: 0.4.3
 Summary: Llama Stack
 Author-email: Meta Llama <llama-oss@meta.com>
 License: MIT
@@ -46,7 +46,7 @@ Requires-Dist: psycopg2-binary
 Requires-Dist: tornado>=6.5.3
 Requires-Dist: urllib3>=2.6.3
 Provides-Extra: client
-Requires-Dist: llama-stack-client==0.4.2; extra == "client"
+Requires-Dist: llama-stack-client==0.4.3; extra == "client"
 Dynamic: license-file
 
 # Llama Stack

{llama_stack-0.4.2.dist-info → llama_stack-0.4.3.dist-info}/RECORD

@@ -30,11 +30,11 @@ llama_stack/core/distribution.py,sha256=pASA0KJ_KTTRQbWP0H5OSvp1ZFQvprbMfE6OLBTq
 llama_stack/core/external.py,sha256=_UTuHkqMzDM07CMAGcPeschNm8NfMkr63iSrLLYc5lg,1869
 llama_stack/core/id_generation.py,sha256=HyTuFALhY0FTpHrYSZbykpPgoNShBeKfFexJz5xEFU8,1185
 llama_stack/core/inspect.py,sha256=0L_Em3kYuvqLQFa4Nn-47rzQtpzhuuWtcnpixCfQuv4,7250
-llama_stack/core/library_client.py,sha256=V5f7apz0heD5DyExwNXiEN0E5xGyQh279BeuVSSzclQ,21380
+llama_stack/core/library_client.py,sha256=Lk1uNUCgnbx8WsCPmtVcwyjwcyU20_CH-hStAR93Fc0,24156
 llama_stack/core/providers.py,sha256=EblMlsWJKGHsXCTmVo-doCJ64JEpBy7-2DoupFkaTUo,5134
 llama_stack/core/request_headers.py,sha256=tUt-RvzUrl7yxbYKBe7nN5YBCgWxShz4cemLvl7XGxc,3692
 llama_stack/core/resolver.py,sha256=IRPPwi60uAe5mlj-NjAR41laP9Dp1WvAI3A-bTMB-mk,19383
-llama_stack/core/stack.py,sha256=YWk2opmFtsYmyEvjrUzDGhLkGV3SNN-omr_eVZKsS-8,27944
+llama_stack/core/stack.py,sha256=cPvwjgldE3L9fEQEReKIIOtHNWEUtUtZx6DL17zIi34,28588
 llama_stack/core/start_stack.sh,sha256=3snlFzur13NS1_UnJQ6t8zK7R5DCRFJKJrz9YTJmWVA,2834
 llama_stack/core/testing_context.py,sha256=TIWetol6Sb2BSiqkq5X0knb0chG03GSpmjByFwVfY60,1438
 llama_stack/core/access_control/__init__.py,sha256=vUvqRS2CXhASaFzYVspRYa5q8usSCzjKUlZhzNLuiKg,200
@@ -54,13 +54,13 @@ llama_stack/core/routers/tool_runtime.py,sha256=bRfPMlLxtdRQ7ad5BPTetw7fi_QT_xV5
 llama_stack/core/routers/vector_io.py,sha256=QqloV8ljTAvjvAytWRBCPnr5Aqm3LidLXalLxDh0W54,21972
 llama_stack/core/routing_tables/__init__.py,sha256=vUvqRS2CXhASaFzYVspRYa5q8usSCzjKUlZhzNLuiKg,200
 llama_stack/core/routing_tables/benchmarks.py,sha256=hxHRQyk8MRVBpdLYIjy4Cim28sdAbPjf9tuZ0msddMg,2552
-llama_stack/core/routing_tables/common.py,sha256=DI7OQvadW9yeV1LVg0FP4d4p0JksbwJVYuVioXzkyQw,10353
+llama_stack/core/routing_tables/common.py,sha256=1jPbybKPYQAllN9HeN5EKxCrLeVHG1kGvRfk2Uy3aQM,10847
 llama_stack/core/routing_tables/datasets.py,sha256=xWhcVUg6TuIcBMRbcYf8_keNc5PADv8PfVEBWmbabJ4,3025
 llama_stack/core/routing_tables/models.py,sha256=PyEYEHlB9uarnSpr6Mnw-bVDOPrANMO06I0dl2orbPg,11388
 llama_stack/core/routing_tables/scoring_functions.py,sha256=gICo9oTTwqT3Ifrbi84wmdcLf2ipiaU3RCoPO6IBt-Y,2464
 llama_stack/core/routing_tables/shields.py,sha256=KjfusJknyZ28evzhdJJLaOuiMJmjkzEPuilIlLH6AcY,2209
 llama_stack/core/routing_tables/toolgroups.py,sha256=X-Uh78mFtfhIpt1RMKHgaam3PD-hWzvmHu0YQj_EEWk,5602
-llama_stack/core/routing_tables/vector_stores.py,sha256=eUvzzWD6Nx9gTFW4ziMR6hGccYRhBKRQYOxbly2knWA,12363
+llama_stack/core/routing_tables/vector_stores.py,sha256=ZLv1SGh_Z2N6eZjCp9_HV5L5nA_b8sio3XeZelP3JdQ,12560
 llama_stack/core/server/__init__.py,sha256=vUvqRS2CXhASaFzYVspRYa5q8usSCzjKUlZhzNLuiKg,200
 llama_stack/core/server/auth.py,sha256=dCyBRg72nplOhKx0xvIiwcptVQmK6140VVCtGjSsvhw,7768
 llama_stack/core/server/auth_providers.py,sha256=1b9eBmvaz2aBQlNkkODaEQHYIF6WejS4bjc3gyr8sg8,20876
@@ -72,19 +72,19 @@ llama_stack/core/storage/__init__.py,sha256=vUvqRS2CXhASaFzYVspRYa5q8usSCzjKUlZh
 llama_stack/core/storage/datatypes.py,sha256=WZqYV_Pgy2nKwEEgjJR1_pST0VUHqY_rnPaWZ8-bjIU,9810
 llama_stack/core/storage/kvstore/__init__.py,sha256=uTrLgzQMAgEUWthwT5LUfHnPAp41Std0zJzlTtPgP-A,309
 llama_stack/core/storage/kvstore/config.py,sha256=XSPhkugrvA9hxBc8DXclD2ktZzNGTQypIbwoH3FXGmI,1436
-llama_stack/core/storage/kvstore/kvstore.py,sha256=Ux5Q8g4fmji4nqKa6QMHQ-gVvhYwn0Jgnmeh0C1Oq8I,4291
+llama_stack/core/storage/kvstore/kvstore.py,sha256=q1LNtQrAxlBBSuVVl6jnyd7v0JHcDpS_7U2h2QmWwtU,4601
 llama_stack/core/storage/kvstore/mongodb/__init__.py,sha256=1LZfkl1Mo3ObnKm7s8619xYiaEZzy95o8pvEifawBQ4,275
-llama_stack/core/storage/kvstore/mongodb/mongodb.py,sha256=oJqt6NKhQZpoaoW0ZSl9iq8WLHmP2SqjguMCaBX2ogk,3453
+llama_stack/core/storage/kvstore/mongodb/mongodb.py,sha256=U8gCCzGhM0KzinwVa9huFzDNZgyVZ4TH7RmHCpH5RJc,3579
 llama_stack/core/storage/kvstore/postgres/__init__.py,sha256=JPP4RZZRUMKOnh9hsnB1mxQkftVfLvo9l2QcXZlUHbY,262
-llama_stack/core/storage/kvstore/postgres/postgres.py,sha256=-d3NQy5DBScMVswYG2lv-p5p74oErOxxdlJlQBsWDQI,4435
+llama_stack/core/storage/kvstore/postgres/postgres.py,sha256=ohEcChgHBjBvkOzIBBtV7Kp1N8kL9CzWRSdSW2-zmMw,4648
 llama_stack/core/storage/kvstore/redis/__init__.py,sha256=FGi8DzDTHkWJufWWoJAXW2zjs2Ji5nsha2cGF84HZis,251
-llama_stack/core/storage/kvstore/redis/redis.py,sha256=6eFkcmVjFCVsSYZ9IGhSWBihzS2UMZhZYKuuoIDPR7A,3517
+llama_stack/core/storage/kvstore/redis/redis.py,sha256=pA41GajiVle-8rbSq1SNExcIdexTpSnk_1vE6zydlfM,3649
 llama_stack/core/storage/kvstore/sqlite/__init__.py,sha256=tnKA5tuANI8w1wIJhhyTYpV5NUxyl5XypjdiHnr0PaY,253
 llama_stack/core/storage/kvstore/sqlite/sqlite.py,sha256=Aecv8Jk8XPOYB_yoU_XkjYUfxqqv2pO7xhVp6F_Urdg,7234
 llama_stack/core/storage/sqlstore/__init__.py,sha256=m-8Cg5fNJnxWZw2KVAHlkYPpGuMcrQliGWiYBQChl3Y,493
 llama_stack/core/storage/sqlstore/authorized_sqlstore.py,sha256=-HWmTl_rQx_97DhFQPx1mREXJr8X3IyplqrpELdEgVc,14474
-llama_stack/core/storage/sqlstore/sqlalchemy_sqlstore.py,sha256=wcqAyT5Bzl5FzdgP7lzGEdsEBG__lbACyw6QzFTfd68,14905
-llama_stack/core/storage/sqlstore/sqlstore.py,sha256=Q3mi-Piv9e8trY6O5z5cQ-izbwcELCLSuWFoR-Y0T98,2980
+llama_stack/core/storage/sqlstore/sqlalchemy_sqlstore.py,sha256=J3pNnZUpUoU3IhPmQpHbLMK7BXguPnCKjQZCYR_wayQ,15248
+llama_stack/core/storage/sqlstore/sqlstore.py,sha256=rGuYGZerCrZZ_G8vMbeP5_ef4jnMFAcXFv-8jniNjno,3228
 llama_stack/core/store/__init__.py,sha256=6rFevukvm0CQxaf3uV3lr5vJTZRB-igWXK2Bd32SNsk,244
 llama_stack/core/store/registry.py,sha256=QSoJYCZ3VjHN3NkqECYGwvv0FJpozuh_9k22Lt_y2tY,7377
 llama_stack/core/utils/__init__.py,sha256=vUvqRS2CXhASaFzYVspRYa5q8usSCzjKUlZhzNLuiKg,200
@@ -200,9 +200,9 @@ llama_stack/providers/inline/agents/meta_reference/agents.py,sha256=Q1fqMRZ2Yfb4
 llama_stack/providers/inline/agents/meta_reference/config.py,sha256=OY82xSuMfh8bMDfV7BZO4P-MMzzjUaMxOayNn4hNPj4,1375
 llama_stack/providers/inline/agents/meta_reference/safety.py,sha256=e5taz2bA9Ks3nlRgLp0CGHsgU_Rv_6ed4bGemQlGEv8,1684
 llama_stack/providers/inline/agents/meta_reference/responses/__init__.py,sha256=vUvqRS2CXhASaFzYVspRYa5q8usSCzjKUlZhzNLuiKg,200
-llama_stack/providers/inline/agents/meta_reference/responses/openai_responses.py,sha256=uLC2Ht9PAsVygAs_p7jRAHxPzm0ZmWFb9k1sU04tNq8,23695
-llama_stack/providers/inline/agents/meta_reference/responses/streaming.py,sha256=PahZY5hfQYJPkuHPe9GKz0FyGawn-AS51zz4LF4CktI,70087
-llama_stack/providers/inline/agents/meta_reference/responses/tool_executor.py,sha256=ga42XVUEQNzvjTLvzvBQpRleU-mvfklUdt8dg08Nid4,22632
+llama_stack/providers/inline/agents/meta_reference/responses/openai_responses.py,sha256=TPuSd7g6sA95b6-oR4ailQ_v6Te7HNqpcVzLFwEFwtU,25192
+llama_stack/providers/inline/agents/meta_reference/responses/streaming.py,sha256=PsHByA2JUBU2KESWXpHq4rBfV4wMBwZ4JMSLZR_kxns,70310
+llama_stack/providers/inline/agents/meta_reference/responses/tool_executor.py,sha256=2aEeqhWP8kCEubFwj9zFj_QzfXD1MssHDedhjvuMZLs,23202
 llama_stack/providers/inline/agents/meta_reference/responses/types.py,sha256=VFUUSbaH15cHzy72pYcxDfuEdKbUxT3F3jg5X-42kbE,8721
 llama_stack/providers/inline/agents/meta_reference/responses/utils.py,sha256=yUGkUZcitTmb6mePDK8Zvw2FDB5FIvxVCx0b887Gqiw,22977
 llama_stack/providers/inline/batches/__init__.py,sha256=vUvqRS2CXhASaFzYVspRYa5q8usSCzjKUlZhzNLuiKg,200
@@ -464,7 +464,7 @@ llama_stack/providers/remote/vector_io/milvus/config.py,sha256=ZRf00ifVSEC5te8Ww
 llama_stack/providers/remote/vector_io/milvus/milvus.py,sha256=epkA43GUBu6u5-uGw1Dk-MNcETTjve4GCGo52p46wnY,16406
 llama_stack/providers/remote/vector_io/pgvector/__init__.py,sha256=yAFgSkT7gwkRtO-xWtm4nkP99QQxZAxsynDUEybIf9Y,564
 llama_stack/providers/remote/vector_io/pgvector/config.py,sha256=yMVFQf4fJDev8ShnbnOicXnLF9r45MIxCKYcbmDFyfw,1540
-llama_stack/providers/remote/vector_io/pgvector/pgvector.py,sha256=UwGoaSFgPOJGp6jdWdMcB5ZSGpeP7RP13TTUfk1wD2U,19601
+llama_stack/providers/remote/vector_io/pgvector/pgvector.py,sha256=LD9SsKHexAg438rBo6Ro479aAWDb1U5o_zTObjr3GI8,20040
 llama_stack/providers/remote/vector_io/qdrant/__init__.py,sha256=Tim4k1BhOROl5BCHuS1YzH6_MilgNdNrv1qe4zFcla4,554
 llama_stack/providers/remote/vector_io/qdrant/config.py,sha256=4akCc4YbYYIBCENj8NRUm3OI-ix0zTOmbgq8RG3nIWU,1115
 llama_stack/providers/remote/vector_io/qdrant/qdrant.py,sha256=VlA-y7F52LC4paHEV6BRQyxWAdBBzh0gWH1hUUs7JMQ,15404
@@ -485,7 +485,7 @@ llama_stack/providers/utils/datasetio/url_utils.py,sha256=Jhxw-bRs5PbkG-pCSRKqoK
 llama_stack/providers/utils/files/__init__.py,sha256=vUvqRS2CXhASaFzYVspRYa5q8usSCzjKUlZhzNLuiKg,200
 llama_stack/providers/utils/files/form_data.py,sha256=-yRXfeSf3AfyRShz6HiSlTPFSa1y4qdlRDxwQOLMK40,2284
 llama_stack/providers/utils/inference/__init__.py,sha256=Ocwqyn7ytwdt1vMFXsPBoa5D6uhA1fIljF-HiIsVvKw,1089
-llama_stack/providers/utils/inference/embedding_mixin.py,sha256=CQ2vvLM1Y6FJcYXdWO5zxVjHud43-W4CjqqdTM47kns,3325
+llama_stack/providers/utils/inference/embedding_mixin.py,sha256=n94bHl0YItWOpRhrCstnz0tr-ZABEJ3DTEGOyWD4t_4,3592
 llama_stack/providers/utils/inference/inference_store.py,sha256=DG_W3FCwA6fjpz9GGko_ylm-4-kNBkxltcXNKOW52OI,9944
 llama_stack/providers/utils/inference/litellm_openai_mixin.py,sha256=MDxQi5KZ1fIjxdpDhvcA7_9Rvu3HNrSQWTTqLrOkUII,13621
 llama_stack/providers/utils/inference/model_registry.py,sha256=CEanQgkbESbAxwczpXQnPHjSa9pGzQBq43tF8t7S9kk,8157
@@ -495,7 +495,7 @@ llama_stack/providers/utils/inference/prompt_adapter.py,sha256=qE2y1EpY0zfYAN6We
 llama_stack/providers/utils/inference/stream_utils.py,sha256=WdM3SPMh9xfOAcpd67_Ld0AaNKrvoYIdJ4nrFOTMmp8,675
 llama_stack/providers/utils/memory/__init__.py,sha256=pA4yikPZUO-A0K2nscz5tEp1yYSBtvglbgC5pe-FGKE,214
 llama_stack/providers/utils/memory/file_utils.py,sha256=MsjispuPO0cMXmRqAoTJ-dwM9uzgYn4aiRFBM-aHP9w,712
-llama_stack/providers/utils/memory/openai_vector_store_mixin.py,sha256=8nq_Nj_pLpznlf7YbsZAHeSucH1FaHD3IILfNoquwpo,58620
+llama_stack/providers/utils/memory/openai_vector_store_mixin.py,sha256=7yVvuq_uc9DxOd7SLqcJOVmu2c0AkBGmPmXZDaM1TgQ,60115
 llama_stack/providers/utils/memory/vector_store.py,sha256=HiNCtZ4OzvPk4RGuRNks7CnbrDoppYyrJdjST-emnZQ,11941
 llama_stack/providers/utils/responses/__init__.py,sha256=vUvqRS2CXhASaFzYVspRYa5q8usSCzjKUlZhzNLuiKg,200
 llama_stack/providers/utils/responses/responses_store.py,sha256=4ziPwlqxMS-mrlB2rL2M8LU9sYewmInH9zN5WPHK00U,10397
@@ -504,7 +504,7 @@ llama_stack/providers/utils/scoring/aggregation_utils.py,sha256=S2XR7DLXEVf7FCTO
 llama_stack/providers/utils/scoring/base_scoring_fn.py,sha256=Upd-Hu8ftoZArfAuHx00uDkvEtmVKWg8U9tgkArro4o,4098
 llama_stack/providers/utils/scoring/basic_scoring_utils.py,sha256=JmGA65N55raHR7rmcdWdTQPaZy4X7I69KFDvfN6716A,714
 llama_stack/providers/utils/tools/__init__.py,sha256=vUvqRS2CXhASaFzYVspRYa5q8usSCzjKUlZhzNLuiKg,200
-llama_stack/providers/utils/tools/mcp.py,sha256=zjGAkmPRRcNALfp9uWaTNKl0WRaCTf7jrASY-Ka1-SA,9315
+llama_stack/providers/utils/tools/mcp.py,sha256=0P0VQgrx7WWoomrmi-1QZFSMlb8DTDLxx8u8TJgMH84,20244
 llama_stack/providers/utils/tools/ttl_dict.py,sha256=4Bv3Nri9HM2FSckfaJJbqzICpO2S_yOXcsgVj_yvsoA,2021
 llama_stack/providers/utils/vector_io/__init__.py,sha256=fGP7xUTCZ3E77v3FtEuGyW2k3S5Tn9x0Kk1aEIafoxA,552
 llama_stack/providers/utils/vector_io/vector_utils.py,sha256=l1asZcxbtlRIaZUi_LbXagclCAveD-k6w28LfOZwqBk,7147
@@ -513,7 +513,7 @@ llama_stack/telemetry/constants.py,sha256=LtXE61xwNL3cBYZXKcXcbwD_Uh1jazP3V8a0od
 llama_stack/telemetry/helpers.py,sha256=7uarMIHL5ngOUXQZxkH96corFxE7Jk5JaizRQ8Z8Ok0,1694
 llama_stack/testing/__init__.py,sha256=vUvqRS2CXhASaFzYVspRYa5q8usSCzjKUlZhzNLuiKg,200
 llama_stack/testing/api_recorder.py,sha256=oGGTrzzBYNNvOIcvcFZenNPthr0yziJ7hlGPtckx460,39240
-llama_stack-0.4.2.dist-info/licenses/LICENSE,sha256=42g1gBn9gHYdBt5e6e1aFYhnc-JT9trU9qBD84oUAlY,1087
+llama_stack-0.4.3.dist-info/licenses/LICENSE,sha256=42g1gBn9gHYdBt5e6e1aFYhnc-JT9trU9qBD84oUAlY,1087
 llama_stack_api/__init__.py,sha256=5XNQGpundjXTutLgnYp6B1t6KITWXH_of626GciNma4,28103
 llama_stack_api/agents.py,sha256=u0sg3AoWCip5o8T4DMTM8uqP3BsdbkKbor3PmxKTg0g,7143
 llama_stack_api/connectors.py,sha256=PcAwndbVQC6pm5HGSlNprqYFTZzhCM7SYHPyRkSIoaQ,4644
@@ -537,7 +537,7 @@ llama_stack_api/scoring_functions.py,sha256=0lP_ZENUh12i12ibg-_XNNPKLHi_TvB8H5Ly
 llama_stack_api/shields.py,sha256=9dNMyTVL0xcR8_BXCHb_zuAJC7Cz8pX8htRwW2-EDSw,2823
 llama_stack_api/tools.py,sha256=eCyZx806VfpBJgsuJF9R3urA8ljF3g0kLapNpx9YRzY,7518
 llama_stack_api/vector_io.py,sha256=3tYy8xLhVvx_rMtfi5Pxv0GwTMm1TfMYwq82tFqRz1U,36517
-llama_stack_api/vector_stores.py,sha256=DMkwPSg05VJOvBJrVlwFU6EHBQEmarADhIzzgt1jjwE,1709
+llama_stack_api/vector_stores.py,sha256=mILSO3k2X-Hg4G3YEdq54fKAenCuAzRAXqpNg-_D_Ng,1832
 llama_stack_api/version.py,sha256=V3jdW3iFPdfOt4jWzJA-di7v0zHLYsn11hNtRzkY7uQ,297
 llama_stack_api/admin/__init__.py,sha256=VnJn9fbk-dFkRrm1P5UWlAOcZDA2jf6dx9W5nt-WgOY,1049
 llama_stack_api/admin/api.py,sha256=m14f4iBUJf-G0qITj66o-TFKCSUiD9U12XRnZ1Slr_w,1961
@@ -575,8 +575,8 @@ llama_stack_api/inspect_api/api.py,sha256=XkdM7jJ3_UlEIE4woEVi5mO2O1aNn9_FPtb18N
 llama_stack_api/inspect_api/fastapi_routes.py,sha256=I7R8roy6einYDzrPN8wNjrRokpoSNZi9zrtmLHS1vDw,2575
 llama_stack_api/inspect_api/models.py,sha256=EW69EHkOG8i0GS8KW8Kz6WaPZV74hzwad8dGXWrrKhs,683
 llama_stack_api/internal/__init__.py,sha256=hZiF7mONpu54guvMUTW9XpfkETUO55u6hqYOYkz8Bt0,307
-llama_stack_api/internal/kvstore.py,sha256=J_lFhhlFcg9uCyn6J758qWSbMIW5nvcfvB66kkitF8g,790
-llama_stack_api/internal/sqlstore.py,sha256=IMOmHiNpxrjqvYNmcsdxbGDUdnMvviFo8AlmT9P27IQ,2219
+llama_stack_api/internal/kvstore.py,sha256=mgNJz6r8_ju3I3JT2Pz5fSX_9DLv_OupsS2NnJe3usY,833
+llama_stack_api/internal/sqlstore.py,sha256=FBIQhG7VOVMMSTe24uMigfxEWXnarY0hzx9HjrNXVnI,2262
 llama_stack_api/llama_stack_api/__init__.py,sha256=5XNQGpundjXTutLgnYp6B1t6KITWXH_of626GciNma4,28103
 llama_stack_api/llama_stack_api/agents.py,sha256=u0sg3AoWCip5o8T4DMTM8uqP3BsdbkKbor3PmxKTg0g,7143
 llama_stack_api/llama_stack_api/connectors.py,sha256=PcAwndbVQC6pm5HGSlNprqYFTZzhCM7SYHPyRkSIoaQ,4644
@@ -600,7 +600,7 @@ llama_stack_api/llama_stack_api/scoring_functions.py,sha256=0lP_ZENUh12i12ibg-_X
 llama_stack_api/llama_stack_api/shields.py,sha256=9dNMyTVL0xcR8_BXCHb_zuAJC7Cz8pX8htRwW2-EDSw,2823
 llama_stack_api/llama_stack_api/tools.py,sha256=eCyZx806VfpBJgsuJF9R3urA8ljF3g0kLapNpx9YRzY,7518
 llama_stack_api/llama_stack_api/vector_io.py,sha256=3tYy8xLhVvx_rMtfi5Pxv0GwTMm1TfMYwq82tFqRz1U,36517
-llama_stack_api/llama_stack_api/vector_stores.py,sha256=DMkwPSg05VJOvBJrVlwFU6EHBQEmarADhIzzgt1jjwE,1709
+llama_stack_api/llama_stack_api/vector_stores.py,sha256=mILSO3k2X-Hg4G3YEdq54fKAenCuAzRAXqpNg-_D_Ng,1832
 llama_stack_api/llama_stack_api/version.py,sha256=V3jdW3iFPdfOt4jWzJA-di7v0zHLYsn11hNtRzkY7uQ,297
 llama_stack_api/llama_stack_api/admin/__init__.py,sha256=VnJn9fbk-dFkRrm1P5UWlAOcZDA2jf6dx9W5nt-WgOY,1049
 llama_stack_api/llama_stack_api/admin/api.py,sha256=m14f4iBUJf-G0qITj66o-TFKCSUiD9U12XRnZ1Slr_w,1961
@@ -638,8 +638,8 @@ llama_stack_api/llama_stack_api/inspect_api/api.py,sha256=XkdM7jJ3_UlEIE4woEVi5m
 llama_stack_api/llama_stack_api/inspect_api/fastapi_routes.py,sha256=I7R8roy6einYDzrPN8wNjrRokpoSNZi9zrtmLHS1vDw,2575
 llama_stack_api/llama_stack_api/inspect_api/models.py,sha256=EW69EHkOG8i0GS8KW8Kz6WaPZV74hzwad8dGXWrrKhs,683
 llama_stack_api/llama_stack_api/internal/__init__.py,sha256=hZiF7mONpu54guvMUTW9XpfkETUO55u6hqYOYkz8Bt0,307
-llama_stack_api/llama_stack_api/internal/kvstore.py,sha256=J_lFhhlFcg9uCyn6J758qWSbMIW5nvcfvB66kkitF8g,790
-llama_stack_api/llama_stack_api/internal/sqlstore.py,sha256=IMOmHiNpxrjqvYNmcsdxbGDUdnMvviFo8AlmT9P27IQ,2219
+llama_stack_api/llama_stack_api/internal/kvstore.py,sha256=mgNJz6r8_ju3I3JT2Pz5fSX_9DLv_OupsS2NnJe3usY,833
+llama_stack_api/llama_stack_api/internal/sqlstore.py,sha256=FBIQhG7VOVMMSTe24uMigfxEWXnarY0hzx9HjrNXVnI,2262
 llama_stack_api/llama_stack_api/providers/__init__.py,sha256=a_187ghsdPNYJ5xLizqKYREJJLBa-lpcIhLp8spgsH8,841
 llama_stack_api/llama_stack_api/providers/api.py,sha256=ytwxri9s6p8j9ClFKgN9mfa1TF0VZh1o8W5cVZR49rc,534
 llama_stack_api/llama_stack_api/providers/fastapi_routes.py,sha256=jb1yrXEk1MdtcgWCToSZtaB-wjKqv5uVKIkvduXoKlM,1962
@@ -648,8 +648,8 @@ llama_stack_api/providers/__init__.py,sha256=a_187ghsdPNYJ5xLizqKYREJJLBa-lpcIhL
 llama_stack_api/providers/api.py,sha256=ytwxri9s6p8j9ClFKgN9mfa1TF0VZh1o8W5cVZR49rc,534
 llama_stack_api/providers/fastapi_routes.py,sha256=jb1yrXEk1MdtcgWCToSZtaB-wjKqv5uVKIkvduXoKlM,1962
 llama_stack_api/providers/models.py,sha256=nqBzh9je_dou35XFjYGD43hwKgjWy6HIRmGWUrcGqOw,653
-llama_stack-0.4.2.dist-info/METADATA,sha256=otXCJCHPmU0tTV9dG5BsYWmxgbWqrlfeicQgZx5OlCc,12464
-llama_stack-0.4.2.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-llama_stack-0.4.2.dist-info/entry_points.txt,sha256=E5xoyAM9064aW_y96eSSwZCNT_ANctrvrhLMJnMQlw0,141
-llama_stack-0.4.2.dist-info/top_level.txt,sha256=pyNYneZU5w62BaExic-GC1ph5kk8JI2mJFwzqiZy2cU,28
-llama_stack-0.4.2.dist-info/RECORD,,
+llama_stack-0.4.3.dist-info/METADATA,sha256=i3L_soHo8hFMex4qrg30BHHC4f79XoW3jdR3Zza76Yk,12464
+llama_stack-0.4.3.dist-info/WHEEL,sha256=wUyA8OaulRlbfwMtmQsvNngGrxQHAvkKcvRmdizlJi0,92
+llama_stack-0.4.3.dist-info/entry_points.txt,sha256=E5xoyAM9064aW_y96eSSwZCNT_ANctrvrhLMJnMQlw0,141
+llama_stack-0.4.3.dist-info/top_level.txt,sha256=pyNYneZU5w62BaExic-GC1ph5kk8JI2mJFwzqiZy2cU,28
+llama_stack-0.4.3.dist-info/RECORD,,

{llama_stack-0.4.2.dist-info → llama_stack-0.4.3.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (80.9.0)
+Generator: setuptools (80.10.2)
 Root-Is-Purelib: true
 Tag: py3-none-any
 

llama_stack_api/internal/kvstore.py

@@ -22,5 +22,7 @@ class KVStore(Protocol):
 
     async def keys_in_range(self, start_key: str, end_key: str) -> list[str]: ...
 
+    async def shutdown(self) -> None: ...
+
 
 __all__ = ["KVStore"]

llama_stack_api/internal/sqlstore.py

@@ -75,5 +75,7 @@ class SqlStore(Protocol):
         nullable: bool = True,
     ) -> None: ...
 
+    async def shutdown(self) -> None: ...
+
 
 __all__ = ["ColumnDefinition", "ColumnType", "SqlStore"]

llama_stack_api/llama_stack_api/internal/kvstore.py

@@ -22,5 +22,7 @@ class KVStore(Protocol):
 
     async def keys_in_range(self, start_key: str, end_key: str) -> list[str]: ...
 
+    async def shutdown(self) -> None: ...
+
 
 __all__ = ["KVStore"]

llama_stack_api/llama_stack_api/internal/sqlstore.py

@@ -75,5 +75,7 @@ class SqlStore(Protocol):
         nullable: bool = True,
     ) -> None: ...
 
+    async def shutdown(self) -> None: ...
+
 
 __all__ = ["ColumnDefinition", "ColumnType", "SqlStore"]

llama_stack_api/llama_stack_api/vector_stores.py

@@ -42,6 +42,7 @@ class VectorStoreInput(BaseModel):
     :param embedding_model: Name of the embedding model to use for vector generation
     :param embedding_dimension: Dimension of the embedding vectors
     :param provider_vector_store_id: (Optional) Provider-specific identifier for the vector store
+    :param vector_store_name: (Optional) Human-readable name for the vector store
     """
 
     vector_store_id: str
@@ -49,3 +50,4 @@ class VectorStoreInput(BaseModel):
     embedding_dimension: int
     provider_id: str | None = None
     provider_vector_store_id: str | None = None
+    vector_store_name: str | None = None

llama_stack_api/vector_stores.py

@@ -42,6 +42,7 @@ class VectorStoreInput(BaseModel):
     :param embedding_model: Name of the embedding model to use for vector generation
     :param embedding_dimension: Dimension of the embedding vectors
     :param provider_vector_store_id: (Optional) Provider-specific identifier for the vector store
+    :param vector_store_name: (Optional) Human-readable name for the vector store
     """
 
     vector_store_id: str
@@ -49,3 +50,4 @@ class VectorStoreInput(BaseModel):
     embedding_dimension: int
     provider_id: str | None = None
     provider_vector_store_id: str | None = None
+    vector_store_name: str | None = None