langgraph-api 0.5.4__py3-none-any.whl → 0.7.3__py3-none-any.whl

langgraph_api/errors.py

@@ -17,8 +17,19 @@ async def http_exception_handler(request: Request, exc: HTTPException) -> Respon
     headers = getattr(exc, "headers", None)
     if not is_body_allowed_for_status_code(exc.status_code):
         return Response(status_code=exc.status_code, headers=headers)
+
+    detail = exc.detail
+    if not detail or not isinstance(detail, str):
+        logger.warning(
+            "HTTPException detail is not a string or was not set",
+            detail_type=type(detail).__name__,
+            status_code=exc.status_code,
+        )
+        # Use safe fallback that won't fail or leak sensitive info
+        detail = "unknown error"
+
     return JSONResponse(
-        {"detail": exc.detail}, status_code=exc.status_code, headers=headers
+        {"detail": detail}, status_code=exc.status_code, headers=headers
     )
 
 

langgraph_api/executor_entrypoint.py

@@ -4,13 +4,10 @@ import json
 import logging.config
 import pathlib
 
-from langgraph_api.queue_entrypoint import main
+from langgraph_api.queue_entrypoint import main as queue_main
 
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser()
 
-    parser.add_argument("--grpc-port", type=int, default=50051)
-    args = parser.parse_args()
+async def main(grpc_port: int = 50051):
     with open(pathlib.Path(__file__).parent.parent / "logging.json") as file:
         loaded_config = json.load(file)
         logging.config.dictConfig(loaded_config)
@@ -23,4 +20,12 @@ if __name__ == "__main__":
     from langgraph_api import config
 
     config.IS_EXECUTOR_ENTRYPOINT = True
-    asyncio.run(main(grpc_port=args.grpc_port, entrypoint_name="python-executor"))
+    await queue_main(grpc_port=grpc_port, entrypoint_name="python-executor")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument("--grpc-port", type=int, default=50051)
+    args = parser.parse_args()
+    asyncio.run(main(grpc_port=args.grpc_port))

langgraph_api/feature_flags.py

@@ -16,3 +16,22 @@ FF_USE_CORE_API = os.getenv("FF_USE_CORE_API", "false").lower() in (
     "1",
     "yes",
 )
+
+# Runtime edition detection
+_RUNTIME_EDITION = os.getenv("LANGGRAPH_RUNTIME_EDITION", "inmem")
+IS_POSTGRES_BACKEND = _RUNTIME_EDITION == "postgres"
+IS_POSTGRES_OR_GRPC_BACKEND = IS_POSTGRES_BACKEND or FF_USE_CORE_API
+# Feature flag for using the JS native API
+FF_USE_JS_API = os.getenv("FF_USE_JS_API", "false").lower() in (
+    "true",
+    "1",
+    "yes",
+)
+
+# In langgraph <= 1.0.3, we automatically subscribed to updates stream events to surface interrupts. In langgraph 1.0.4 we include interrupts in values events (which we are automatically subscribed to), so we no longer need to implicitly subscribe to updates stream events
+# If the version is not valid, e.g. rc/alpha/etc., we default to 0.0.0
+try:
+    LANGGRAPH_PY_PATCH = tuple(map(int, __version__.split(".")[:3]))
+except ValueError:
+    LANGGRAPH_PY_PATCH = (0, 0, 0)
+UPDATES_NEEDED_FOR_INTERRUPTS = LANGGRAPH_PY_PATCH <= (1, 0, 3)

langgraph_api/graph.py

@@ -3,17 +3,20 @@ import functools
 import glob
 import importlib.util
 import inspect
+import logging
 import os
 import sys
+import time
 import warnings
 from collections.abc import AsyncIterator, Callable
 from contextlib import asynccontextmanager
 from itertools import filterfalse
-from typing import TYPE_CHECKING, Any, NamedTuple, TypeGuard, cast
+from typing import Any, NamedTuple, TypeGuard, cast
 from uuid import UUID, uuid5
 
 import orjson
 import structlog
+from langchain_core.embeddings import Embeddings  # noqa: TC002
 from langgraph.checkpoint.base import BaseCheckpointSaver
 from langgraph.constants import CONFIG_KEY_CHECKPOINTER
 from langgraph.graph import StateGraph
@@ -22,15 +25,17 @@ from langgraph.store.base import BaseStore
 from starlette.exceptions import HTTPException
 
 from langgraph_api import config as lg_api_config
-from langgraph_api.feature_flags import FF_USE_CORE_API, USE_RUNTIME_CONTEXT_API
+from langgraph_api import timing
+from langgraph_api.feature_flags import (
+    IS_POSTGRES_OR_GRPC_BACKEND,
+    USE_RUNTIME_CONTEXT_API,
+)
 from langgraph_api.js.base import BaseRemotePregel, is_js_path
 from langgraph_api.schema import Config
+from langgraph_api.timing import profiled_import
 from langgraph_api.utils.config import run_in_executor, var_child_runnable_config
 from langgraph_api.utils.errors import GraphLoadError
 
-if TYPE_CHECKING:
-    from langchain_core.embeddings import Embeddings
-
 logger = structlog.stdlib.get_logger(__name__)
 
 GraphFactoryFromConfig = Callable[[Config], Pregel | StateGraph]
@@ -51,13 +56,13 @@ async def register_graph(
     description: str | None = None,
 ) -> None:
     """Register a graph."""
-    from langgraph_api.grpc_ops.ops import Assistants as AssistantsGrpc
     from langgraph_runtime.database import connect
-    from langgraph_runtime.ops import Assistants as AssistantsRuntime
 
-    Assistants = AssistantsGrpc if FF_USE_CORE_API else AssistantsRuntime
+    if IS_POSTGRES_OR_GRPC_BACKEND:
+        from langgraph_api.grpc.ops import Assistants
+    else:
+        from langgraph_runtime.ops import Assistants
 
-    await logger.ainfo(f"Registering graph with id '{graph_id}'", graph_id=graph_id)
     GRAPHS[graph_id] = graph
     if callable(graph):
         FACTORY_ACCEPTS_CONFIG[graph_id] = len(inspect.signature(graph).parameters) > 0
@@ -91,19 +96,70 @@ async def register_graph(
         await register_graph_db()
 
 
+def _validate_assistant_id(assistant_id: str) -> None:
+    """Validate an assistant ID is either a graph_id or a valid UUID. Throw an error if not valid."""
+    if assistant_id and assistant_id not in GRAPHS:
+        # Not a graph_id, must be a valid UUID
+        try:
+            UUID(assistant_id)
+        except ValueError:
+            # Invalid format - return 404 to match test expectations
+            raise HTTPException(
+                status_code=404,
+                detail=f"Assistant '{assistant_id}' not found",
+            ) from None
+
+
+def _log_slow_graph_generation(
+    start: float,
+    value_type: str,
+    graph_id: str,
+    warn_threshold_ms: float = 100,
+    error_threshold_ms: float = 250,
+) -> None:
+    """Log warning/error if graph generation was slow."""
+    elapsed_secs = time.perf_counter() - start
+    elapsed_ms = elapsed_secs * 1000
+    elapsed_ms_rounded = round(elapsed_ms, 2)
+    log_level = None
+    if elapsed_ms > error_threshold_ms:
+        log_level = logging.ERROR
+    elif elapsed_ms > warn_threshold_ms:
+        log_level = logging.WARNING
+    if log_level is not None:
+        logger.log(
+            log_level,
+            f"Slow graph load. Accessing graph '{graph_id}' took {elapsed_ms_rounded}ms."
+            " Move expensive initialization (API clients, DB connections, model loading)"
+            " from graph factory if you are seeing API slowness.",
+            elapsed_ms=elapsed_ms_rounded,
+            value_type=value_type,
+            graph_id=graph_id,
+        )
+
+
 @asynccontextmanager
-async def _generate_graph(value: Any) -> AsyncIterator[Any]:
-    """Yield a graph object regardless of its type."""
+async def _generate_graph(value: Any, graph_id: str) -> AsyncIterator[Any]:
+    """Yield a graph object regardless of its type.
+
+    Logs a warning if graph generation takes >100ms, error if >250ms.
+    """
+    start = time.perf_counter()
+    value_type = type(value).__name__
     if isinstance(value, Pregel | BaseRemotePregel):
         yield value
     elif hasattr(value, "__aenter__") and hasattr(value, "__aexit__"):
         async with value as ctx_value:
+            _log_slow_graph_generation(start, value_type, graph_id)
             yield ctx_value
     elif hasattr(value, "__enter__") and hasattr(value, "__exit__"):
         with value as ctx_value:
+            _log_slow_graph_generation(start, value_type, graph_id)
             yield ctx_value
     elif asyncio.iscoroutine(value):
-        yield await value
+        result = await value
+        _log_slow_graph_generation(start, value_type, graph_id)
+        yield result
     else:
         yield value
 
@@ -132,14 +188,18 @@ async def get_graph(
     *,
     checkpointer: BaseCheckpointSaver | None = None,
     store: BaseStore | None = None,
+    is_for_execution: bool = True,
 ) -> AsyncIterator[Pregel]:
     """Return the runnable."""
     from langgraph_api.utils import config as lg_config
+    from langgraph_api.utils import merge_auth
 
     assert_graph_exists(graph_id)
     value = GRAPHS[graph_id]
     if is_factory(value, graph_id):
         config = lg_config.ensure_config(config)
+        config["configurable"]["__is_for_execution__"] = is_for_execution
+        config = merge_auth(config)
 
         if store is not None:
             if USE_RUNTIME_CONTEXT_API:
@@ -152,7 +212,7 @@ async def get_graph(
                 elif isinstance(runtime, dict):
                     patched_runtime = Runtime(**(runtime | {"store": store}))
                 elif runtime.store is None:
-                    patched_runtime = cast(Runtime, runtime).override(store=store)
+                    patched_runtime = cast("Runtime", runtime).override(store=store)
                 else:
                     patched_runtime = runtime
 
@@ -170,7 +230,7 @@ async def get_graph(
         var_child_runnable_config.set(config)
         value = value(config) if factory_accepts_config(value, graph_id) else value()
     try:
-        async with _generate_graph(value) as graph_obj:
+        async with _generate_graph(value, graph_id) as graph_obj:
             if isinstance(graph_obj, StateGraph):
                 graph_obj = graph_obj.compile()
             if not isinstance(graph_obj, Pregel | BaseRemotePregel):
@@ -232,9 +292,9 @@ class GraphSpec(NamedTuple):
     variable: str | None = None
     config: dict | None = None
     """The configuration for the graph.
-    
+
     Contains information such as: tags, recursion_limit and configurable.
-    
+
     Configurable is a dict containing user defined values for the graph.
     """
     description: str | None = None
@@ -429,7 +489,7 @@ async def collect_graphs_from_env(register: bool = False) -> None:
 def _handle_exception(task: asyncio.Task) -> None:
     try:
         task.result()
-    except asyncio.CancelledError:
+    except (asyncio.CancelledError, SystemExit):
         pass
     except Exception as e:
         logger.exception("Task failed", exc_info=e)
@@ -448,42 +508,59 @@ def verify_graphs() -> None:
     asyncio.run(collect_graphs_from_env())
 
 
+def _metadata_fn(spec: GraphSpec) -> dict[str, Any]:
+    return {"graph_id": spec.id, "module": spec.module, "path": spec.path}
+
+
+@timing.timer(
+    message="Importing graph with id {graph_id}",
+    metadata_fn=_metadata_fn,
+    warn_threshold_secs=3,
+    warn_message=(
+        "Import for graph {graph_id} exceeded the expected startup time. "
+        "Slow initialization (often due to work executed at import time) can delay readiness, "
+        "reduce scale-out capacity, and may cause deployments to be marked unhealthy."
+    ),
+    error_threshold_secs=30,
+)
 def _graph_from_spec(spec: GraphSpec) -> GraphValue:
     """Return a graph from a spec."""
     # import the graph module
-    if spec.module:
-        module = importlib.import_module(spec.module)
-    elif spec.path:
-        try:
-            modname = (
-                spec.path.replace("/", "__")
-                .replace(".py", "")
-                .replace(" ", "_")
-                .lstrip(".")
-            )
-            modspec = importlib.util.spec_from_file_location(modname, spec.path)
-            if modspec is None:
-                raise ValueError(f"Could not find python file for graph: {spec}")
-            module = importlib.util.module_from_spec(modspec)
-            sys.modules[modname] = module
-            modspec.loader.exec_module(module)  # type: ignore[possibly-unbound-attribute]
-        except ImportError as e:
-            e.add_note(f"Could not import python module for graph:\n{spec}")
-            if lg_api_config.API_VARIANT == "local_dev":
-                e.add_note(
-                    "This error likely means you haven't installed your project and its dependencies yet. Before running the server, install your project:\n\n"
-                    "If you are using requirements.txt:\n"
-                    "python -m pip install -r requirements.txt\n\n"
-                    "If you are using pyproject.toml or setuptools:\n"
-                    "python -m pip install -e .\n\n"
-                    "Make sure to run this command from your project's root directory (where your setup.py or pyproject.toml is located)"
+    import_path = f"{spec.module or spec.path}:{spec.variable or '<auto>'}"
+    with profiled_import(import_path):
+        if spec.module:
+            module = importlib.import_module(spec.module)
+        elif spec.path:
+            try:
+                modname = (
+                    spec.path.replace("/", "__")
+                    .replace(".py", "")
+                    .replace(" ", "_")
+                    .lstrip(".")
                 )
-            raise
-        except FileNotFoundError as e:
-            e.add_note(f"Could not find python file for graph: {spec}")
-            raise
-    else:
-        raise ValueError("Graph specification must have a path or module")
+                modspec = importlib.util.spec_from_file_location(modname, spec.path)
+                if modspec is None:
+                    raise ValueError(f"Could not find python file for graph: {spec}")
+                module = importlib.util.module_from_spec(modspec)
+                sys.modules[modname] = module
+                modspec.loader.exec_module(module)  # type: ignore[possibly-unbound-attribute]
+            except ImportError as e:
+                e.add_note(f"Could not import python module for graph:\n{spec}")
+                if lg_api_config.API_VARIANT == "local_dev":
+                    e.add_note(
+                        "This error likely means you haven't installed your project and its dependencies yet. Before running the server, install your project:\n\n"
+                        "If you are using requirements.txt:\n"
+                        "python -m pip install -r requirements.txt\n\n"
+                        "If you are using pyproject.toml or setuptools:\n"
+                        "python -m pip install -e .\n\n"
+                        "Make sure to run this command from your project's root directory (where your setup.py or pyproject.toml is located)"
+                    )
+                raise
+            except FileNotFoundError as e:
+                e.add_note(f"Could not find python file for graph: {spec}")
+                raise
+        else:
+            raise ValueError("Graph specification must have a path or module")
 
     if spec.variable:
         try:
@@ -589,6 +666,13 @@ def _get_init_embeddings() -> Callable[[str, ...], "Embeddings"] | None:
         return None
 
 
+@timing.timer(
+    message="Loading embeddings {embeddings_path}",
+    metadata_fn=lambda index_config: {"embeddings_path": index_config.get("embed")},
+    warn_threshold_secs=5,
+    warn_message="Loading embeddings '{embeddings_path}' took longer than expected",
+    error_threshold_secs=10,
+)
 def resolve_embeddings(index_config: dict) -> "Embeddings":
     """Return embeddings from config.
 
@@ -607,26 +691,41 @@ def resolve_embeddings(index_config: dict) -> "Embeddings":
     from langchain_core.embeddings import Embeddings
     from langgraph.store.base import ensure_embeddings
 
-    embed: str = index_config["embed"]
+    embed = index_config["embed"]
+    if isinstance(embed, Embeddings):
+        return embed
+    if callable(embed):
+        return ensure_embeddings(embed)
+    if not isinstance(embed, str):
+        raise ValueError(
+            f"Embeddings config must be a string or callable, got: {type(embed).__name__}"
+        )
     if ".py:" in embed:
         module_name, function = embed.rsplit(":", 1)
         module_name = module_name.rstrip(":")
 
         try:
-            if "/" in module_name:
-                # Load from file path
-                modname = (
-                    module_name.replace("/", "__").replace(".py", "").replace(" ", "_")
-                )
-                modspec = importlib.util.spec_from_file_location(modname, module_name)
-                if modspec is None:
-                    raise ValueError(f"Could not find embeddings file: {module_name}")
-                module = importlib.util.module_from_spec(modspec)
-                sys.modules[modname] = module
-                modspec.loader.exec_module(module)  # type: ignore[possibly-unbound-attribute]
-            else:
-                # Load from Python module
-                module = importlib.import_module(module_name)
+            with profiled_import(embed):
+                if "/" in module_name:
+                    # Load from file path
+                    modname = (
+                        module_name.replace("/", "__")
+                        .replace(".py", "")
+                        .replace(" ", "_")
+                    )
+                    modspec = importlib.util.spec_from_file_location(
+                        modname, module_name
+                    )
+                    if modspec is None:
+                        raise ValueError(
+                            f"Could not find embeddings file: {module_name}"
+                        )
+                    module = importlib.util.module_from_spec(modspec)
+                    sys.modules[modname] = module
+                    modspec.loader.exec_module(module)  # type: ignore[possibly-unbound-attribute]
+                else:
+                    # Load from Python module
+                    module = importlib.import_module(module_name)
 
             embedding_fn = getattr(module, function, None)
             if embedding_fn is None:

langgraph_api/{grpc_ops → grpc}/client.py

@@ -1,18 +1,34 @@
 """gRPC client wrapper for LangGraph persistence services."""
 
 import asyncio
-import os
+import threading
+import time
 
 import structlog
 from grpc import aio  # type: ignore[import]
+from grpc_health.v1 import health_pb2, health_pb2_grpc  # type: ignore[import]
 
-from .generated.core_api_pb2_grpc import AdminStub, AssistantsStub, ThreadsStub
+from langgraph_api import config
+
+from .generated.checkpointer_pb2_grpc import CheckpointerStub
+from .generated.core_api_pb2_grpc import (
+    AdminStub,
+    AssistantsStub,
+    RunsStub,
+    ThreadsStub,
+)
 
 logger = structlog.stdlib.get_logger(__name__)
 
 
-# Shared global client pool
+# Shared gRPC client pools (main thread + thread-local for isolated loops).
 _client_pool: "GrpcClientPool | None" = None
+_thread_local = threading.local()
+
+
+GRPC_HEALTHCHECK_TIMEOUT = 5.0
+GRPC_INIT_TIMEOUT = 10.0
+GRPC_INIT_PROBE_INTERVAL = 0.5
 
 
 class GrpcClient:
@@ -27,13 +43,14 @@ class GrpcClient:
         Args:
             server_address: The gRPC server address (default: localhost:50051)
         """
-        self.server_address = server_address or os.getenv(
-            "GRPC_SERVER_ADDRESS", "localhost:50051"
-        )
+        self.server_address = server_address or config.GRPC_SERVER_ADDRESS
         self._channel: aio.Channel | None = None
         self._assistants_stub: AssistantsStub | None = None
+        self._runs_stub: RunsStub | None = None
         self._threads_stub: ThreadsStub | None = None
         self._admin_stub: AdminStub | None = None
+        self._checkpointer_stub: CheckpointerStub | None = None
+        self._health_stub: health_pb2_grpc.HealthStub | None = None
 
     async def __aenter__(self):
         """Async context manager entry."""
@@ -49,11 +66,19 @@ class GrpcClient:
         if self._channel is not None:
             return
 
-        self._channel = aio.insecure_channel(self.server_address)
+        options = [
+            ("grpc.max_receive_message_length", config.GRPC_CLIENT_MAX_RECV_MSG_BYTES),
+            ("grpc.max_send_message_length", config.GRPC_CLIENT_MAX_SEND_MSG_BYTES),
+        ]
+
+        self._channel = aio.insecure_channel(self.server_address, options=options)
 
         self._assistants_stub = AssistantsStub(self._channel)
+        self._runs_stub = RunsStub(self._channel)
         self._threads_stub = ThreadsStub(self._channel)
         self._admin_stub = AdminStub(self._channel)
+        self._checkpointer_stub = CheckpointerStub(self._channel)
+        self._health_stub = health_pb2_grpc.HealthStub(self._channel)
 
         await logger.adebug(
             "Connected to gRPC server", server_address=self.server_address
@@ -65,10 +90,37 @@ class GrpcClient:
             await self._channel.close()
             self._channel = None
             self._assistants_stub = None
+            self._runs_stub = None
             self._threads_stub = None
             self._admin_stub = None
+            self._checkpointer_stub = None
+            self._health_stub = None
             await logger.adebug("Closed gRPC connection")
 
+    async def healthcheck(self) -> bool:
+        """Check if the gRPC server is healthy.
+
+        Returns:
+            True if the server is healthy and serving.
+
+        Raises:
+            RuntimeError: If the client is not connected or the server is unhealthy.
+        """
+        if self._health_stub is None:
+            raise RuntimeError(
+                "Client not connected. Use async context manager or call connect() first."
+            )
+
+        request = health_pb2.HealthCheckRequest(service="")
+        response = await self._health_stub.Check(
+            request, timeout=GRPC_HEALTHCHECK_TIMEOUT
+        )
+
+        if response.status != health_pb2.HealthCheckResponse.SERVING:
+            raise RuntimeError(f"gRPC server is not healthy. Status: {response.status}")
+
+        return True
+
     @property
     def assistants(self) -> AssistantsStub:
         """Get the assistants service stub."""
@@ -87,6 +139,15 @@ class GrpcClient:
             )
         return self._threads_stub
 
+    @property
+    def runs(self) -> RunsStub:
+        """Get the runs service stub."""
+        if self._runs_stub is None:
+            raise RuntimeError(
+                "Client not connected. Use async context manager or call connect() first."
+            )
+        return self._runs_stub
+
     @property
     def admin(self) -> AdminStub:
         """Get the admin service stub."""
@@ -96,6 +157,15 @@ class GrpcClient:
             )
         return self._admin_stub
 
+    @property
+    def checkpointer(self) -> CheckpointerStub:
+        """Get the checkpointer service stub."""
+        if self._checkpointer_stub is None:
+            raise RuntimeError(
+                "Client not connected. Use async context manager or call connect() first."
+            )
+        return self._checkpointer_stub
+
 
 class GrpcClientPool:
     """Pool of gRPC clients for load distribution."""
@@ -158,25 +228,85 @@ async def get_shared_client() -> GrpcClient:
 
     Uses a pool of channels for better performance under high concurrency.
     Each channel is a separate TCP connection that can handle ~100-200
-    concurrent streams effectively.
+    concurrent streams effectively. Pools are scoped per thread/loop to
+    avoid cross-loop gRPC channel usage.
 
     Returns:
         A GrpcClient instance from the pool
     """
+    if threading.current_thread() is not threading.main_thread():
+        pool = getattr(_thread_local, "grpc_pool", None)
+        if pool is None:
+            pool = GrpcClientPool(
+                pool_size=1,
+                server_address=config.GRPC_SERVER_ADDRESS,
+            )
+            _thread_local.grpc_pool = pool
+        return await pool.get_client()
+
     global _client_pool
     if _client_pool is None:
-        from langgraph_api import config
-
         _client_pool = GrpcClientPool(
             pool_size=config.GRPC_CLIENT_POOL_SIZE,
-            server_address=os.getenv("GRPC_SERVER_ADDRESS"),
+            server_address=config.GRPC_SERVER_ADDRESS,
         )
-
     return await _client_pool.get_client()
 
 
+async def wait_until_grpc_ready(
+    timeout_seconds: float = GRPC_INIT_TIMEOUT,
+    interval_seconds: float = GRPC_INIT_PROBE_INTERVAL,
+):
+    """Wait for the gRPC server to be ready with retries during startup.
+
+    Args:
+        timeout_seconds: Maximum time to wait for the server to be ready.
+        interval_seconds: Time to wait between health check attempts.
+    Raises:
+        RuntimeError: If the server is not ready within the timeout period.
+    """
+    client = await get_shared_client()
+    max_attempts = int(timeout_seconds / interval_seconds)
+
+    await logger.ainfo(
+        "Waiting for gRPC server to be ready",
+        timeout_seconds=timeout_seconds,
+        interval_seconds=interval_seconds,
+        max_attempts=max_attempts,
+    )
+    start_time = time.time()
+    for attempt in range(max_attempts):
+        try:
+            await client.healthcheck()
+            await logger.ainfo(
+                "gRPC server is ready",
+                attempt=attempt + 1,
+                elapsed_seconds=round(time.time() - start_time, 3),
+            )
+            return
+        except Exception as exc:
+            if attempt >= max_attempts - 1:
+                raise RuntimeError(
+                    f"gRPC server not ready after {timeout_seconds}s (reached max attempts: {max_attempts})"
+                ) from exc
+            else:
+                await logger.adebug(
+                    "Waiting for gRPC server to be ready",
+                    attempt=attempt + 1,
+                    max_attempts=max_attempts,
+                )
+                await asyncio.sleep(interval_seconds)
+
+
 async def close_shared_client():
     """Close the shared gRPC client pool."""
+    if threading.current_thread() is not threading.main_thread():
+        pool = getattr(_thread_local, "grpc_pool", None)
+        if pool is not None:
+            await pool.close()
+            delattr(_thread_local, "grpc_pool")
+        return
+
     global _client_pool
     if _client_pool is not None:
         await _client_pool.close()