eval-protocol 0.0.3__py3-none-any.whl

eval_protocol/mcp_agent/intermediary_server.py

@@ -0,0 +1,542 @@
+import asyncio
+import logging
+import uuid
+from typing import Any, Dict, List, Optional
+
+import anyio  # Added for debugging cancel scopes and tasks
+from mcp import types as mcp_types  # Added for type hinting
+from pydantic import BaseModel, Field
+
+from eval_protocol.mcp_agent.config import AppConfig, BackendServerConfig
+from eval_protocol.mcp_agent.orchestration.base_client import (
+    AbstractOrchestrationClient,
+    ManagedInstanceInfo,
+)
+from eval_protocol.mcp_agent.orchestration.local_docker_client import (
+    LocalDockerOrchestrationClient,
+)
+from eval_protocol.mcp_agent.orchestration.remote_http_client import (
+    RemoteHttpOrchestrationClient,
+)
+from eval_protocol.mcp_agent.session import IntermediarySessionData
+
+logger = logging.getLogger(__name__)
+# logger.setLevel(logging.DEBUG)  # Removed: Let level be set by main config
+
+from mcp.server.fastmcp.server import Context as FastMCPContext
+from mcp.server.fastmcp.server import FastMCP
+
+# RequestContext is not directly used by handlers anymore, mcp_ctx is.
+
+
+# Backend initialization models (moved here to avoid separate backends module)
+class BackendInitRequest(BaseModel):
+    backend_name_ref: str = Field(
+        ...,
+        description="The unique reference name of the backend configuration to use (must match one in AppConfig.backends).",
+    )
+    num_instances: int = Field(
+        1,
+        ge=1,
+        description="Number of instances of this backend to provision for the session.",
+    )
+    template_details: Optional[Any] = Field(
+        None,
+        description="Backend-specific details for initializing stateful instances from a template.",
+    )
+
+    class Config:
+        extra = "forbid"
+
+
+class BackendInitResult(BaseModel):
+    backend_name_ref: str
+    instances: List[ManagedInstanceInfo]
+
+
+# Pydantic models for tool arguments
+class InitializeSessionArgs(BaseModel):
+    backends: List[BackendInitRequest]
+
+
+class CallBackendToolArgs(BaseModel):
+    rk_session_id: str = Field(..., description="The session ID obtained from initialize_session.")
+    backend_name_ref: str = Field(..., description="The reference name of the backend to target.")
+    instance_id: str = Field(..., description="The ID of the specific backend instance to use.")
+    tool_name: str = Field(..., description="The name of the tool to call on the backend instance.")
+    tool_args: Dict[str, Any] = Field(default_factory=dict, description="Arguments for the backend tool.")
+
+
+class ListBackendToolsArgs(BaseModel):
+    rk_session_id: str = Field(..., description="The session ID obtained from initialize_session.")
+    backend_name_ref: str = Field(..., description="The reference name of the backend to target.")
+    instance_id: str = Field(..., description="The ID of the specific backend instance to query for tools.")
+
+
+class CleanupSessionArgs(BaseModel):
+    rk_session_id: str = Field(..., description="The session ID to clean up.")
+
+
+# Ping might not need specific args if it uses session from mcp_ctx, or could take rk_session_id
+class PingArgs(BaseModel):
+    rk_session_id: Optional[str] = Field(default=None, description="Optional session ID for context.")
+
+
+class RewardKitIntermediaryServer(FastMCP):
+    def __init__(self, app_config: AppConfig, **kwargs_for_fastmcp):
+        super().__init__(
+            name="RewardKitIntermediaryMCP",
+            instructions="Intermediary Server for managing backend MCP resources for RewardKit RL rollouts.",
+            **kwargs_for_fastmcp,
+        )
+
+        self.app_config = app_config
+        self._local_docker_orchestrator: Optional[LocalDockerOrchestrationClient] = None
+        self._remote_http_orchestrators: Dict[str, RemoteHttpOrchestrationClient] = {}
+        self._shared_global_instances: Dict[str, ManagedInstanceInfo] = {}
+        self._shared_instance_locks: Dict[str, asyncio.Lock] = {}
+        self.intermediary_session_data: Dict[str, IntermediarySessionData] = {}
+
+        logger.info("RewardKitIntermediaryServer (FastMCP based) initialized. AppConfig loaded.")
+
+        # Register tools directly
+        self.add_tool(self._initialize_session_actual, name="initialize_session")
+        self.add_tool(self._call_backend_tool_actual, name="call_backend_tool")
+        self.add_tool(self._list_backend_tools_actual, name="list_backend_tools")  # New tool
+        self.add_tool(self._cleanup_session_actual, name="cleanup_session")
+        self.add_tool(self._ping_actual, name="ping")
+
+        logger.info("Registered tools directly with FastMCP.")
+
+        # Explicitly set this module's logger level based on app_config
+        # This is to ensure it overrides any prior default or hardcoded DEBUG level
+        # if external configuration in main.py isn't fully effective.
+        try:
+            config_log_level_str = app_config.log_level.upper()
+            config_log_level_int = getattr(logging, config_log_level_str, logging.INFO)
+            if logger.getEffectiveLevel() != config_log_level_int:
+                logger.info(
+                    f"Overriding intermediary_server logger level from {logging.getLevelName(logger.getEffectiveLevel())} to {config_log_level_str}"
+                )
+                logger.setLevel(config_log_level_int)
+                # Also ensure handlers attached directly to this logger respect it (if any)
+                for handler in logger.handlers:
+                    handler.setLevel(config_log_level_int)
+            logger.info(
+                f"IntermediaryServer logger effective level: {logging.getLevelName(logger.getEffectiveLevel())}"
+            )
+
+        except Exception as e_log:
+            logger.error(f"Error trying to set intermediary_server logger level: {e_log}")
+
+    # Removed _execute_proxied_tool_impl and _internal_tool_handlers
+
+    async def _initialize_orchestrators(self):
+        logger.info("Initializing orchestration clients...")
+        if any(b.orchestration_mode == "local_docker" for b in self.app_config.backends):
+            self._local_docker_orchestrator = LocalDockerOrchestrationClient(self.app_config)
+            await self._local_docker_orchestrator.startup()
+            logger.info("LocalDockerOrchestrationClient initialized and started.")
+
+        unique_remote_api_refs = set()
+        for backend_cfg in self.app_config.backends:
+            if backend_cfg.orchestration_mode == "remote_http_api":
+                if backend_cfg.remote_api_config_ref:
+                    unique_remote_api_refs.add(backend_cfg.remote_api_config_ref)
+                elif backend_cfg.remote_api_config_inline:
+                    logger.warning(
+                        f"Inline remote_api_config for {backend_cfg.backend_name_ref}. Consider using global_remote_apis."
+                    )
+                    key = backend_cfg.remote_api_config_inline.base_url
+                    if key not in self._remote_http_orchestrators:
+                        temp_app_config_for_inline = AppConfig(
+                            global_remote_apis={key: backend_cfg.remote_api_config_inline}
+                        )
+                        client = RemoteHttpOrchestrationClient(temp_app_config_for_inline)
+                        await client.startup()
+                        self._remote_http_orchestrators[key] = client
+                        logger.info(f"RemoteHttpOrchestrationClient for inline config {key} initialized.")
+
+        for ref_name in unique_remote_api_refs:
+            if ref_name not in self.app_config.global_remote_apis:
+                logger.error(f"Remote API ref '{ref_name}' not in global_remote_apis.")
+                continue
+            if ref_name not in self._remote_http_orchestrators:
+                isolated_app_cfg = AppConfig(
+                    global_remote_apis={ref_name: self.app_config.global_remote_apis[ref_name]},
+                    global_remote_api_defaults=self.app_config.global_remote_api_defaults,
+                )
+                client = RemoteHttpOrchestrationClient(isolated_app_cfg)
+                await client.startup()
+                self._remote_http_orchestrators[ref_name] = client
+                logger.info(f"RemoteHttpOrchestrationClient for '{ref_name}' initialized.")
+        logger.info("Orchestration clients initialization complete.")
+
+    def _get_orchestration_client(self, backend_cfg: BackendServerConfig) -> AbstractOrchestrationClient:
+        if backend_cfg.orchestration_mode == "local_docker":
+            if not self._local_docker_orchestrator:
+                raise RuntimeError("Local Docker orchestrator not initialized.")
+            return self._local_docker_orchestrator
+        elif backend_cfg.orchestration_mode == "remote_http_api":
+            key = backend_cfg.remote_api_config_ref
+            if not key:
+                if backend_cfg.remote_api_config_inline:
+                    key = backend_cfg.remote_api_config_inline.base_url
+                else:
+                    raise ValueError(f"Remote API config missing for {backend_cfg.backend_name_ref}")
+            client = self._remote_http_orchestrators.get(key)
+            if not client:
+                raise RuntimeError(f"Remote HTTP orchestrator for '{key}' not initialized.")
+            return client
+        else:
+            raise ValueError(f"Unsupported orchestration mode: {backend_cfg.orchestration_mode}")
+
+    async def _get_or_provision_shared_global_instance(self, backend_name_ref: str) -> ManagedInstanceInfo:
+        if backend_name_ref not in self._shared_instance_locks:
+            self._shared_instance_locks[backend_name_ref] = asyncio.Lock()
+        async with self._shared_instance_locks[backend_name_ref]:
+            if backend_name_ref in self._shared_global_instances:
+                logger.info(f"Returning existing shared global instance for '{backend_name_ref}'.")
+                return self._shared_global_instances[backend_name_ref]
+            logger.info(f"Provisioning new shared global instance for '{backend_name_ref}'.")
+            backend_cfg = next(
+                (b for b in self.app_config.backends if b.backend_name_ref == backend_name_ref),
+                None,
+            )
+            if not backend_cfg or backend_cfg.instance_scoping != "shared_global":
+                raise ValueError(f"Backend '{backend_name_ref}' not for shared_global scoping.")
+            orchestration_client = self._get_orchestration_client(backend_cfg)
+            provisioned_list = await orchestration_client.provision_instances(
+                backend_config=backend_cfg,
+                num_instances=1,
+                session_id="global_shared_session",
+                template_details=backend_cfg.template_data_path_host,
+            )
+            if not provisioned_list:
+                raise RuntimeError(f"Failed to provision shared global for '{backend_name_ref}'.")
+            instance_info = provisioned_list[0]
+            self._shared_global_instances[backend_name_ref] = instance_info
+            logger.info(f"Provisioned shared global for '{backend_name_ref}': {instance_info.instance_id}")
+            return instance_info
+
+    async def _provision_shared_global_instances(self):
+        logger.info("Pre-provisioning all shared_global instances...")
+        for backend_cfg in self.app_config.backends:
+            if backend_cfg.instance_scoping == "shared_global":
+                try:
+                    await self._get_or_provision_shared_global_instance(backend_cfg.backend_name_ref)
+                except Exception as e:
+                    logger.error(
+                        f"Failed to pre-provision for '{backend_cfg.backend_name_ref}': {e}",
+                        exc_info=True,
+                    )
+        logger.info("Shared_global instances pre-provisioning complete.")
+
+    async def _initialize_session_actual(self, mcp_ctx: FastMCPContext, args: InitializeSessionArgs) -> Dict[str, Any]:
+        task_name = anyio.get_current_task().name if anyio.get_current_task() else "unknown_task"
+        logger.debug(
+            f"ENTERING _initialize_session_actual: task='{task_name}', mcp_ctx type: {type(mcp_ctx)}, args: {args}"
+        )
+
+        transport_session_id: Optional[str] = None
+        if (
+            hasattr(mcp_ctx, "session")
+            and mcp_ctx.session
+            and hasattr(mcp_ctx.session, "client_params")
+            and mcp_ctx.session.client_params
+            and hasattr(mcp_ctx.session.client_params, "session_id")
+            and mcp_ctx.session.client_params.session_id
+        ):
+            transport_session_id = mcp_ctx.session.client_params.session_id
+            logger.info(f"Retrieved transport_session_id: {transport_session_id}")
+
+        rk_session_id = transport_session_id if transport_session_id else uuid.uuid4().hex
+        if not transport_session_id:
+            logger.warning(f"Transport session ID not found. Generated new rk_session_id: {rk_session_id}")
+        else:
+            logger.info(f"Using transport_session_id as rk_session_id: {rk_session_id}")
+
+        if rk_session_id in self.intermediary_session_data:
+            logger.warning(f"rk_session_id '{rk_session_id}' already exists. Overwriting.")
+        session_data = IntermediarySessionData(session_id=rk_session_id)
+        self.intermediary_session_data[rk_session_id] = session_data
+
+        logger.info(
+            f"Initializing IntermediarySessionData for rk_session_id '{rk_session_id}' with {len(args.backends)} backend requests."
+        )
+        initialized_backends_results: List[BackendInitResult] = []
+
+        for backend_req in args.backends:
+            backend_cfg = next(
+                (b for b in self.app_config.backends if b.backend_name_ref == backend_req.backend_name_ref),
+                None,
+            )
+            if not backend_cfg:
+                logger.error(f"Session {rk_session_id}: Config for '{backend_req.backend_name_ref}' not found.")
+                initialized_backends_results.append(
+                    BackendInitResult(backend_name_ref=backend_req.backend_name_ref, instances=[])
+                )
+                continue
+            try:
+                if backend_cfg.instance_scoping == "shared_global":
+                    shared_instance_info = await self._get_or_provision_shared_global_instance(
+                        backend_req.backend_name_ref
+                    )
+                    instances_for_this_backend = [shared_instance_info] * backend_req.num_instances
+                else:
+                    orchestration_client = self._get_orchestration_client(backend_cfg)
+                    instances_for_this_backend = await orchestration_client.provision_instances(
+                        backend_config=backend_cfg,
+                        num_instances=backend_req.num_instances,
+                        session_id=session_data.session_id,
+                        template_details=backend_req.template_details,
+                    )
+                session_data.add_managed_instances(backend_req.backend_name_ref, instances_for_this_backend)
+                initialized_backends_results.append(
+                    BackendInitResult(
+                        backend_name_ref=backend_req.backend_name_ref,
+                        instances=instances_for_this_backend,
+                    )
+                )
+            except Exception as e:
+                logger.error(
+                    f"Session {rk_session_id}: Error initializing '{backend_req.backend_name_ref}': {e}",
+                    exc_info=True,
+                )
+                initialized_backends_results.append(
+                    BackendInitResult(
+                        backend_name_ref=backend_req.backend_name_ref,
+                        instances=[],
+                        error_message=str(e),
+                    )
+                )
+
+        task_name_exit = anyio.get_current_task().name if anyio.get_current_task() else "unknown_task"
+        logger.debug(f"EXITING _initialize_session_actual: task='{task_name_exit}'")
+        return {
+            "rk_session_id": rk_session_id,
+            "initialized_backends": [res.model_dump(exclude_none=True) for res in initialized_backends_results],
+        }
+
+    async def _call_backend_tool_actual(self, mcp_ctx: FastMCPContext, args: CallBackendToolArgs) -> Dict[str, Any]:
+        task_name_entry = anyio.get_current_task().name if anyio.get_current_task() else "unknown_task"
+        logger.debug(
+            f"ENTERING _call_backend_tool_actual: task='{task_name_entry}', mcp_ctx type: {type(mcp_ctx)}, args: {args}"
+        )
+
+        session_data = self.intermediary_session_data.get(args.rk_session_id)
+        if not session_data:
+            task_name_error = anyio.get_current_task().name if anyio.get_current_task() else "unknown_task"
+            logger.error(
+                f"ERROR in _call_backend_tool_actual (session not found): task='{task_name_error}', rk_session_id='{args.rk_session_id}'"
+            )
+            raise ValueError(f"IntermediarySessionData for rk_session_id '{args.rk_session_id}' not found.")
+
+        target_instances = session_data.get_managed_instances(args.backend_name_ref, args.instance_id)
+        if not target_instances:
+            raise ValueError(
+                f"Instance '{args.instance_id}' for backend '{args.backend_name_ref}' not found in session '{args.rk_session_id}'."
+            )
+        managed_instance_info = target_instances[0]
+        backend_cfg = next(
+            (b for b in self.app_config.backends if b.backend_name_ref == args.backend_name_ref),
+            None,
+        )
+        if not backend_cfg:
+            raise ValueError(f"Backend config '{args.backend_name_ref}' not found.")
+        orchestration_client = self._get_orchestration_client(backend_cfg)
+
+        task_name_before_call = anyio.get_current_task().name if anyio.get_current_task() else "unknown_task"
+        logger.debug(
+            f"BEFORE orchestrator.call_tool_on_instance in _call_backend_tool_actual: task='{task_name_before_call}'"
+        )
+
+        try:
+            result = await orchestration_client.call_tool_on_instance(
+                instance=managed_instance_info,
+                tool_name=args.tool_name,
+                tool_args=args.tool_args,
+            )
+            task_name_after_call = anyio.get_current_task().name if anyio.get_current_task() else "unknown_task"
+            logger.debug(
+                f"AFTER orchestrator.call_tool_on_instance in _call_backend_tool_actual: task='{task_name_after_call}'"
+            )
+
+            task_name_exit = anyio.get_current_task().name if anyio.get_current_task() else "unknown_task"
+            logger.debug(f"EXITING _call_backend_tool_actual (SUCCESS): task='{task_name_exit}'")
+            return result
+        except Exception as e:
+            task_name_exception = anyio.get_current_task().name if anyio.get_current_task() else "unknown_task"
+            logger.error(
+                f"EXCEPTION in _call_backend_tool_actual: task='{task_name_exception}'. Session {args.rk_session_id}: Error calling tool '{args.tool_name}' on instance '{args.instance_id}': {e}",
+                exc_info=True,
+            )
+            raise
+
+    async def _list_backend_tools_actual(
+        self, mcp_ctx: FastMCPContext, args: ListBackendToolsArgs
+    ) -> Dict[str, Any]:  # Returning dict for FastMCP, will be ListToolsResult internally
+        task_name_entry = anyio.get_current_task().name if anyio.get_current_task() else "unknown_task"
+        logger.debug(f"ENTERING _list_backend_tools_actual: task='{task_name_entry}', args: {args}")
+
+        session_data = self.intermediary_session_data.get(args.rk_session_id)
+        if not session_data:
+            logger.error(
+                f"ERROR in _list_backend_tools_actual (session not found): rk_session_id='{args.rk_session_id}'"
+            )
+            raise ValueError(f"IntermediarySessionData for rk_session_id '{args.rk_session_id}' not found.")
+
+        target_instances = session_data.get_managed_instances(args.backend_name_ref, args.instance_id)
+        if not target_instances:
+            raise ValueError(
+                f"Instance '{args.instance_id}' for backend '{args.backend_name_ref}' not found in session '{args.rk_session_id}'."
+            )
+        managed_instance_info = target_instances[0]
+
+        backend_cfg = next(
+            (b for b in self.app_config.backends if b.backend_name_ref == args.backend_name_ref),
+            None,
+        )
+        if not backend_cfg:
+            raise ValueError(f"Backend config '{args.backend_name_ref}' not found.")
+        orchestration_client = self._get_orchestration_client(backend_cfg)
+
+        logger.debug(
+            f"Calling orchestrator.list_tools_on_instance for backend '{args.backend_name_ref}', instance '{args.instance_id}'"
+        )
+        try:
+            list_tools_result: mcp_types.ListToolsResult = await orchestration_client.list_tools_on_instance(
+                instance=managed_instance_info
+            )
+            # FastMCP tools expect to return a dictionary that can be JSON serialized.
+            # ListToolsResult is a Pydantic model, so model_dump() is appropriate.
+            return list_tools_result.model_dump(exclude_none=True)
+        except Exception as e:
+            logger.error(
+                f"EXCEPTION in _list_backend_tools_actual for session {args.rk_session_id}, backend {args.backend_name_ref}, instance {args.instance_id}: {e}",
+                exc_info=True,
+            )
+            raise  # Re-raise to let FastMCP handle error reporting to client
+
+    async def cleanup_session_internal(self, session_data_to_clean: IntermediarySessionData, rk_session_id: str):
+        logger.info(f"Starting internal cleanup for IntermediarySessionData (rk_session_id: '{rk_session_id}').")
+        all_session_instances = session_data_to_clean.get_all_managed_instances()
+        local_docker_instances = [inst for inst in all_session_instances if inst.orchestration_mode == "local_docker"]
+        if local_docker_instances and self._local_docker_orchestrator:
+            try:
+                await self._local_docker_orchestrator.deprovision_instances(local_docker_instances)
+            except Exception as e:
+                logger.error(
+                    f"Session {rk_session_id}: Error deprovisioning local Docker: {e}",
+                    exc_info=True,
+                )
+
+        remote_instances_by_key: Dict[str, List[ManagedInstanceInfo]] = {}
+        for inst in all_session_instances:
+            if inst.orchestration_mode == "remote_http_api":
+                key = self._get_orchestration_client_key_for_instance(inst)
+                if key:
+                    remote_instances_by_key.setdefault(key, []).append(inst)
+        for key, remote_list in remote_instances_by_key.items():
+            orchestrator = self._remote_http_orchestrators.get(key)
+            if orchestrator and remote_list:
+                try:
+                    await orchestrator.deprovision_instances(remote_list)
+                except Exception as e:
+                    logger.error(
+                        f"Session {rk_session_id}: Error deprovisioning remote for '{key}': {e}",
+                        exc_info=True,
+                    )
+        logger.info(f"Internal cleanup for session data (rk_session_id: '{rk_session_id}') complete.")
+
+    async def _cleanup_session_actual(self, mcp_ctx: FastMCPContext, args: CleanupSessionArgs) -> Dict[str, str]:
+        logger.debug(f"_cleanup_session_actual called. mcp_ctx type: {type(mcp_ctx)}, args: {args}")
+        session_data_obj = self.intermediary_session_data.pop(args.rk_session_id, None)
+        if not session_data_obj:
+            logger.warning(
+                f"IntermediarySessionData for rk_session_id '{args.rk_session_id}' not found or already cleaned."
+            )
+            return {
+                "status": "custom_session_data_not_found_or_already_cleaned",
+                "rk_session_id": args.rk_session_id,
+            }
+        await self.cleanup_session_internal(session_data_obj, args.rk_session_id)
+        logger.info(f"IntermediarySessionData for rk_session_id '{args.rk_session_id}' fully cleaned up.")
+        return {"status": "cleaned", "rk_session_id": args.rk_session_id}
+
+    async def startup(self):
+        logger.info("RewardKitIntermediaryServer performing custom startup tasks...")
+        try:
+            await self._initialize_orchestrators()
+            await self._provision_shared_global_instances()
+            logger.info("RewardKitIntermediaryServer custom startup tasks complete.")
+        except Exception as e:
+            logger.error(
+                f"Error during RewardKitIntermediaryServer custom startup: {e}",
+                exc_info=True,
+            )
+            raise
+
+    async def _ping_actual(self, mcp_ctx: FastMCPContext, args: PingArgs) -> Dict[str, str]:
+        logger.debug(f"_ping_actual called. mcp_ctx type: {type(mcp_ctx)}, args: {args}")
+        ping_session_id: Optional[str] = None
+        if args.rk_session_id:  # If client provides its known rk_session_id
+            ping_session_id = args.rk_session_id
+            logger.info(f"Ping using rk_session_id from args: {ping_session_id}")
+        elif (
+            hasattr(mcp_ctx, "session")
+            and mcp_ctx.session
+            and hasattr(mcp_ctx.session, "client_params")
+            and mcp_ctx.session.client_params
+            and hasattr(mcp_ctx.session.client_params, "session_id")
+            and mcp_ctx.session.client_params.session_id
+        ):
+            ping_session_id = mcp_ctx.session.client_params.session_id
+            logger.info(f"Ping using transport_session_id from mcp_ctx: {ping_session_id}")
+        else:
+            ping_session_id = "unknown_session_for_ping"
+            logger.warning(f"Session ID for ping not found in args or mcp_ctx, using fallback: {ping_session_id}")
+        return {"reply": "pong", "session_id": ping_session_id or ""}
+
+    async def shutdown(self):
+        logger.info("RewardKitIntermediaryServer (FastMCP based) performing custom shutdown tasks...")
+        logger.info(f"Cleaning up {len(self.intermediary_session_data)} IntermediarySessionData entries...")
+        for session_id_key in list(self.intermediary_session_data.keys()):
+            session_data_obj = self.intermediary_session_data.pop(session_id_key, None)
+            if session_data_obj:
+                await self.cleanup_session_internal(session_data_obj, session_id_key)
+
+        shared_instances = list(self._shared_global_instances.values())
+        if shared_instances:
+            logger.info(f"Deprovisioning {len(shared_instances)} shared global instances.")
+            local_shared = [i for i in shared_instances if i.orchestration_mode == "local_docker"]
+            if local_shared and self._local_docker_orchestrator:
+                await self._local_docker_orchestrator.deprovision_instances(local_shared)
+            remote_shared_by_key: Dict[str, List[ManagedInstanceInfo]] = {}
+            for inst_info in shared_instances:
+                if inst_info.orchestration_mode == "remote_http_api":
+                    key = self._get_orchestration_client_key_for_instance(inst_info)
+                    if key:
+                        remote_shared_by_key.setdefault(key, []).append(inst_info)
+            for key, instances_list in remote_shared_by_key.items():
+                orchestrator = self._remote_http_orchestrators.get(key)
+                if orchestrator:
+                    await orchestrator.deprovision_instances(instances_list)
+
+        if self._local_docker_orchestrator:
+            await self._local_docker_orchestrator.shutdown()
+        for orch in self._remote_http_orchestrators.values():
+            await orch.shutdown()
+        logger.info("RewardKitIntermediaryServer custom shutdown tasks complete.")
+
+    def _get_orchestration_client_key_for_instance(self, instance_info: ManagedInstanceInfo) -> Optional[str]:
+        if instance_info.orchestration_mode == "remote_http_api":
+            backend_cfg = next(
+                (b for b in self.app_config.backends if b.backend_name_ref == instance_info.backend_name_ref),
+                None,
+            )
+            if backend_cfg:
+                return backend_cfg.remote_api_config_ref or (
+                    backend_cfg.remote_api_config_inline.base_url if backend_cfg.remote_api_config_inline else None
+                )
+        return None