feat(proxy): route custom extensions through sidecar

litellm/extensions/__init__.py

@@ -0,0 +1,5 @@
+"""Client-side adapters for the external Python extension host."""
+
+from .runtime import configure_extension_runtime, get_extension_runtime
+
+__all__ = ("configure_extension_runtime", "get_extension_runtime")

litellm/extensions/adapters.py

@@ -0,0 +1,467 @@
+# pyright: reportAny=false, reportUnknownArgumentType=false, reportUnknownMemberType=false
+# pyright: reportUnknownParameterType=false, reportUnknownVariableType=false
+# pyright: reportMissingParameterType=false
+from __future__ import annotations
+
+import hashlib
+import time
+from collections.abc import AsyncGenerator, AsyncIterator, Mapping
+from datetime import datetime
+from typing import Final
+
+from litellm.caching import DualCache
+from litellm.exceptions import GuardrailRaisedException
+from litellm.integrations.custom_guardrail import CustomGuardrail
+from litellm.integrations.custom_logger import CustomLogger
+from litellm.proxy._types import UserAPIKeyAuth
+from litellm.python_extension.generated.v1 import extension_host_pb2 as pb
+from litellm.python_extension_host.compatibility import decode_json, encode_json, response_from_json
+from litellm.types.guardrails import GuardrailEventHooks, Mode
+from litellm.types.utils import CallTypesLiteral, ModelResponseStream
+
+from .cache_gateway import InvocationCacheRegistry
+from .client import PythonExtensionClient
+
+
+class _RemoteHooks:
+    def __init__(
+        self,
+        client: PythonExtensionClient,
+        plugin_id: str,
+        cache_registry: InvocationCacheRegistry,
+    ) -> None:
+        self._extension_client: Final = client
+        self._extension_plugin_id: Final = plugin_id
+        self._extension_cache_registry: Final = cache_registry
+
+    async def execute(
+        self,
+        phase: pb.HookPhase,
+        data: dict[str, object],  # mutable-ok: LiteLLM compatibility payload
+        auth: UserAPIKeyAuth,
+        call_type: str,
+        cache: object | None = None,
+        response: object | None = None,
+    ) -> object | None:
+        context = _invocation_context(  # rebind-ok: invocation-scoped RPC state
+            data, self._extension_client.manifest.revision_id, call_type
+        )  # rebind-ok: invocation-scoped RPC state
+        cache_ref: Final = self._extension_cache_registry.register(context.invocation_id, cache)
+        invocation: Final = pb.GuardrailInvocation(
+            context=context,
+            plugin_id=self._extension_plugin_id,
+            hook_phase=phase,
+            request_json=encode_json(data),
+            auth=_auth_context(auth),
+        )
+        if response is not None:
+            invocation.response_json = encode_json(response)
+        if cache_ref is not None:
+            invocation.cache.CopyFrom(cache_ref)
+        try:
+            result = await self._extension_client.execute_guardrail(  # rebind-ok: invocation-scoped RPC state
+                invocation
+            )  # rebind-ok: invocation-scoped RPC state
+        finally:
+            self._extension_cache_registry.revoke(cache_ref)
+        if result.decision in (
+            pb.GUARDRAIL_DECISION_ALLOW,
+            pb.GUARDRAIL_DECISION_ERROR,
+            pb.GUARDRAIL_DECISION_UNSPECIFIED,
+        ):
+            return response if phase == pb.HOOK_PHASE_POST_CALL else None
+        if result.decision == pb.GUARDRAIL_DECISION_BLOCK:
+            public_error: Final = result.public_error if result.HasField("public_error") else None
+            raise GuardrailRaisedException(
+                guardrail_name=self._extension_plugin_id,
+                message=public_error.message if public_error is not None else "request blocked by extension",
+                should_wrap_with_default_message=False,
+                status_code=(
+                    public_error.status_code
+                    if public_error is not None and public_error.HasField("status_code")
+                    else 400
+                ),
+                blocked_content=True,
+            )
+        if result.decision == pb.GUARDRAIL_DECISION_REPLACE_REQUEST and result.HasField("request_json"):
+            replacement: Final = decode_json(result.request_json, "request_json")
+            if isinstance(replacement, dict):
+                data.clear()
+                data.update(replacement)
+                return data
+        if result.decision == pb.GUARDRAIL_DECISION_REPLACE_RESPONSE and result.HasField("response_json"):
+            return response_from_json(result.response_json)
+        return response if phase == pb.HOOK_PHASE_POST_CALL else None
+
+    def enqueue_event(
+        self,
+        kind: pb.CallbackEventKind,
+        kwargs: dict[str, object],  # mutable-ok: LiteLLM compatibility payload
+        response: object,
+        start_time: object,
+        end_time: object,
+        streaming: bool = False,
+    ) -> None:
+        context = _invocation_context(  # rebind-ok: invocation-scoped RPC state
+            kwargs,
+            self._extension_client.manifest.revision_id,
+            str(kwargs.get("call_type", "unknown")),
+        )
+        payload: Final = kwargs.get("standard_logging_payload", kwargs)
+        event: Final = pb.CallbackEvent(
+            context=context,
+            plugin_id=self._extension_plugin_id,
+            kind=kind,
+            standard_logging_payload_json=encode_json(payload),
+            start_time_seconds=_timestamp(start_time),
+            end_time_seconds=_timestamp(end_time),
+            auth=_auth_context_from_kwargs(kwargs),
+            streaming=streaming,
+        )
+        if response is not None and not isinstance(response, Exception):
+            event.response_json = encode_json(response)
+        if isinstance(response, Exception):
+            event.error_json = encode_json(
+                {"type": type(response).__name__, "message": str(response)}  # mutable-ok: LiteLLM compatibility payload
+            )  # mutable-ok: LiteLLM compatibility payload
+        self._extension_client.enqueue_callback(event)
+
+    async def transform_iterator(
+        self,
+        auth: UserAPIKeyAuth,
+        response: AsyncIterator[ModelResponseStream],
+        request_data: dict[str, object],  # mutable-ok: LiteLLM compatibility payload
+    ) -> AsyncGenerator[ModelResponseStream, None]:
+        context = _invocation_context(  # rebind-ok: invocation-scoped RPC state
+            request_data,
+            self._extension_client.manifest.revision_id,
+            str(request_data.get("call_type", "stream")),
+        )
+        pending: list[  # mutable-ok: LiteLLM compatibility payload # rebind-ok: invocation-scoped RPC state
+            ModelResponseStream
+        ] = []  # mutable-ok: LiteLLM compatibility payload # rebind-ok: invocation-scoped RPC state
+
+        async def frames() -> AsyncIterator[pb.StreamFrame]:
+            yield pb.StreamFrame(
+                kind=pb.STREAM_FRAME_KIND_OPEN,
+                stream_id=context.invocation_id,
+                open=pb.StreamOpen(
+                    context=context,
+                    plugin_id=self._extension_plugin_id,
+                    request_json=encode_json(request_data),
+                    auth=_auth_context(auth),
+                    iterator_hook=(
+                        "async_post_call_streaming_iterator_hook"
+                        in self._extension_client.descriptor_hooks(self._extension_plugin_id)
+                    ),
+                ),
+            )
+            async for chunk in response:
+                pending.append(chunk)
+                yield pb.StreamFrame(
+                    kind=pb.STREAM_FRAME_KIND_INPUT_CHUNK,
+                    stream_id=context.invocation_id,
+                    chunk_json=encode_json(chunk),
+                )
+            yield pb.StreamFrame(kind=pb.STREAM_FRAME_KIND_END, stream_id=context.invocation_id)
+
+        completed = False  # rebind-ok: invocation-scoped RPC state
+        async for frame in self._extension_client.transform_stream(frames()):
+            if frame.kind == pb.STREAM_FRAME_KIND_OUTPUT_CHUNK and frame.HasField("chunk_json"):
+                transformed = response_from_json(frame.chunk_json)
+                if isinstance(transformed, ModelResponseStream):
+                    yield transformed
+            elif frame.kind == pb.STREAM_FRAME_KIND_END:
+                completed = True
+            elif frame.kind == pb.STREAM_FRAME_KIND_ERROR:
+                break
+        if not completed:
+            for chunk in pending:
+                yield chunk
+            async for chunk in response:
+                yield chunk
+
+
+class RemoteCustomLogger(CustomLogger):
+    def __init__(
+        self,
+        client: PythonExtensionClient,
+        plugin_id: str,
+        cache_registry: InvocationCacheRegistry,
+    ) -> None:
+        super().__init__()
+        self._remote = _RemoteHooks(client, plugin_id, cache_registry)
+
+    async def async_pre_call_hook(
+        self,
+        user_api_key_dict: UserAPIKeyAuth,
+        cache: DualCache,
+        data: dict[str, object],  # mutable-ok: LiteLLM compatibility payload
+        call_type: CallTypesLiteral,
+    ) -> Exception | str | dict[str, object] | None:  # mutable-ok: LiteLLM compatibility payload
+        result = await self._remote.execute(  # rebind-ok: invocation-scoped RPC state
+            pb.HOOK_PHASE_PRE_CALL,
+            data,
+            user_api_key_dict,
+            call_type,
+            cache=cache,
+        )
+        return result if isinstance(result, Exception | str | dict) else None
+
+    async def async_moderation_hook(
+        self,
+        data: dict[str, object],  # mutable-ok: LiteLLM compatibility payload
+        user_api_key_dict: UserAPIKeyAuth,
+        call_type: CallTypesLiteral,
+    ) -> object | None:
+        return await self._remote.execute(pb.HOOK_PHASE_DURING_CALL, data, user_api_key_dict, call_type)
+
+    async def async_post_call_success_hook(
+        self,
+        data: dict[str, object],  # mutable-ok: LiteLLM compatibility payload
+        user_api_key_dict: UserAPIKeyAuth,
+        response: object,
+    ) -> object | None:
+        return await self._remote.execute(
+            pb.HOOK_PHASE_POST_CALL,
+            data,
+            user_api_key_dict,
+            str(data.get("call_type", "unknown")),
+            response=response,
+        )
+
+    async def async_log_success_event(
+        self,
+        kwargs: dict[str, object],  # mutable-ok: LiteLLM compatibility payload
+        response_obj: object,
+        start_time: object,
+        end_time: object,
+    ) -> None:
+        self._remote.enqueue_event(pb.CALLBACK_EVENT_KIND_SUCCESS, kwargs, response_obj, start_time, end_time)
+
+    async def async_log_failure_event(
+        self,
+        kwargs: dict[str, object],  # mutable-ok: LiteLLM compatibility payload
+        response_obj: object,
+        start_time: object,
+        end_time: object,
+    ) -> None:
+        self._remote.enqueue_event(pb.CALLBACK_EVENT_KIND_FAILURE, kwargs, response_obj, start_time, end_time)
+
+    async def async_log_stream_event(
+        self,
+        kwargs: dict[str, object],  # mutable-ok: LiteLLM compatibility payload
+        response_obj: object,
+        start_time: object,
+        end_time: object,
+    ) -> None:
+        self._remote.enqueue_event(
+            pb.CALLBACK_EVENT_KIND_SUCCESS,
+            kwargs,
+            response_obj,
+            start_time,
+            end_time,
+            streaming=True,
+        )
+
+    async def async_post_call_streaming_iterator_hook(
+        self,
+        user_api_key_dict: UserAPIKeyAuth,
+        response: AsyncIterator[ModelResponseStream],
+        request_data: dict[str, object],  # mutable-ok: LiteLLM compatibility payload
+    ) -> AsyncGenerator[ModelResponseStream, None]:
+        async for chunk in self._remote.transform_iterator(user_api_key_dict, response, request_data):
+            yield chunk
+
+
+class RemoteCustomGuardrail(CustomGuardrail):
+    def __init__(
+        self,
+        client: PythonExtensionClient,
+        plugin_id: str,
+        cache_registry: InvocationCacheRegistry,
+        guardrail_name: str | None = None,
+        event_hook: GuardrailEventHooks  # mutable-ok: LiteLLM compatibility payload
+        | list[GuardrailEventHooks]
+        | Mode
+        | None = None,  # mutable-ok: LiteLLM compatibility payload
+        default_on: bool = False,
+        extra_params: Mapping[str, object] | None = None,
+    ) -> None:
+        params = (  # rebind-ok: invocation-scoped RPC state
+            extra_params or {}  # mutable-ok: LiteLLM compatibility payload
+        )  # mutable-ok: LiteLLM compatibility payload # rebind-ok: invocation-scoped RPC state
+        super().__init__(
+            guardrail_name=guardrail_name,
+            event_hook=event_hook,
+            default_on=default_on,
+            mask_request_content=_bool_param(params, "mask_request_content", False),
+            mask_response_content=_bool_param(params, "mask_response_content", False),
+            violation_message_template=_str_param(params, "violation_message_template"),
+            end_session_after_n_fails=_int_param(params, "end_session_after_n_fails"),
+            on_violation=_str_param(params, "on_violation"),
+            realtime_violation_message=_str_param(params, "realtime_violation_message"),
+            on_sensitive_data=_str_param(params, "on_sensitive_data"),
+            sensitive_data_route_to_model=_str_param(params, "sensitive_data_route_to_model"),
+            sticky_session_routing=_bool_param(params, "sticky_session_routing", True),
+            run_in_parallel=_bool_param(params, "run_in_parallel", False),
+            scan_raw_request=_bool_param(params, "scan_raw_request", False),
+            only_scan_new_messages=_bool_param(params, "only_scan_new_messages", False),
+            turn_off_message_logging=_bool_param(params, "turn_off_message_logging", False),
+            message_logging=_bool_param(params, "message_logging", True),
+        )
+        self._remote = _RemoteHooks(client, plugin_id, cache_registry)
+
+    async def async_pre_call_hook(
+        self,
+        user_api_key_dict: UserAPIKeyAuth,
+        cache: DualCache,
+        data: dict[str, object],  # mutable-ok: LiteLLM compatibility payload
+        call_type: CallTypesLiteral,
+    ) -> Exception | str | dict[str, object] | None:  # mutable-ok: LiteLLM compatibility payload
+        result = await self._remote.execute(  # rebind-ok: invocation-scoped RPC state
+            pb.HOOK_PHASE_PRE_CALL,
+            data,
+            user_api_key_dict,
+            call_type,
+            cache=cache,
+        )
+        return result if isinstance(result, Exception | str | dict) else None
+
+    async def async_moderation_hook(
+        self,
+        data: dict[str, object],  # mutable-ok: LiteLLM compatibility payload
+        user_api_key_dict: UserAPIKeyAuth,
+        call_type: CallTypesLiteral,
+    ) -> object | None:
+        return await self._remote.execute(pb.HOOK_PHASE_DURING_CALL, data, user_api_key_dict, call_type)
+
+    async def async_post_call_success_hook(
+        self,
+        data: dict[str, object],  # mutable-ok: LiteLLM compatibility payload
+        user_api_key_dict: UserAPIKeyAuth,
+        response: object,
+    ) -> object | None:
+        return await self._remote.execute(
+            pb.HOOK_PHASE_POST_CALL,
+            data,
+            user_api_key_dict,
+            str(data.get("call_type", "unknown")),
+            response=response,
+        )
+
+    async def async_post_call_streaming_iterator_hook(
+        self,
+        user_api_key_dict: UserAPIKeyAuth,
+        response: AsyncIterator[ModelResponseStream],
+        request_data: dict[str, object],  # mutable-ok: LiteLLM compatibility payload
+    ) -> AsyncGenerator[ModelResponseStream, None]:
+        async for chunk in self._remote.transform_iterator(user_api_key_dict, response, request_data):
+            yield chunk
+
+
+def _invocation_context(data: Mapping[str, object], revision: str, call_type: str) -> pb.InvocationContext:
+    metadata = data.get("metadata")  # rebind-ok: invocation-scoped RPC state
+    metadata_mapping = (  # rebind-ok: invocation-scoped RPC state
+        metadata if isinstance(metadata, Mapping) else {}  # mutable-ok: LiteLLM compatibility payload
+    )  # mutable-ok: LiteLLM compatibility payload # rebind-ok: invocation-scoped RPC state
+    request_id: Final = str(
+        data.get("request_id") or data.get("litellm_call_id") or metadata_mapping.get("request_id") or ""
+    )
+    invocation_id: Final = str(
+        data.get("litellm_call_id") or request_id or hashlib.sha256(str(time.time_ns()).encode()).hexdigest()[:24]
+    )
+    return pb.InvocationContext(
+        request_id=request_id,
+        invocation_id=invocation_id,
+        active_revision=revision,
+        api_surface=str(data.get("api_surface", data.get("call_type", "unknown"))),
+        call_type=call_type,
+        trace_context=_trace_context(metadata_mapping),
+    )
+
+
+def _auth_context(auth: UserAPIKeyAuth) -> pb.AuthContext:
+    key = getattr(auth, "api_key", None) or getattr(auth, "token", None) or ""  # rebind-ok: invocation-scoped RPC state
+    metadata = getattr(auth, "metadata", None)  # rebind-ok: invocation-scoped RPC state
+    return pb.AuthContext(
+        key_hash=_hash_key(str(key)) if key else "",
+        user_id=str(getattr(auth, "user_id", None) or ""),
+        team_id=str(getattr(auth, "team_id", None) or ""),
+        request_metadata=_safe_metadata(
+            metadata if isinstance(metadata, Mapping) else {}  # mutable-ok: LiteLLM compatibility payload
+        ),  # mutable-ok: LiteLLM compatibility payload
+    )
+
+
+def _auth_context_from_kwargs(kwargs: Mapping[str, object]) -> pb.AuthContext:
+    litellm_params: Final = kwargs.get("litellm_params")
+    params = (  # rebind-ok: invocation-scoped RPC state
+        litellm_params if isinstance(litellm_params, Mapping) else {}  # mutable-ok: LiteLLM compatibility payload
+    )  # mutable-ok: LiteLLM compatibility payload # rebind-ok: invocation-scoped RPC state
+    metadata = params.get("metadata")  # rebind-ok: invocation-scoped RPC state
+    metadata_mapping = (  # rebind-ok: invocation-scoped RPC state
+        metadata if isinstance(metadata, Mapping) else {}  # mutable-ok: LiteLLM compatibility payload
+    )  # mutable-ok: LiteLLM compatibility payload # rebind-ok: invocation-scoped RPC state
+    key = params.get("api_key") or metadata_mapping.get("user_api_key") or ""  # rebind-ok: invocation-scoped RPC state
+    return pb.AuthContext(
+        key_hash=_hash_key(str(key)) if key else "",
+        user_id=str(metadata_mapping.get("user_api_key_user_id") or ""),
+        team_id=str(metadata_mapping.get("user_api_key_team_id") or ""),
+        request_metadata=_safe_metadata(metadata_mapping),
+    )
+
+
+def _hash_key(value: str) -> str:
+    return hashlib.sha256(value.encode()).hexdigest()
+
+
+def _safe_metadata(
+    metadata: Mapping[object, object],
+) -> dict[str, str]:  # mutable-ok: LiteLLM compatibility payload
+    denied: Final = (
+        "authorization",
+        "api_key",
+        "token",
+        "cookie",
+        "secret",
+        "password",
+    )
+    return {  # mutable-ok: LiteLLM compatibility payload
+        str(key): str(value)
+        for key, value in metadata.items()
+        if not any(part in str(key).lower() for part in denied) and isinstance(value, str | int | float | bool)
+    }
+
+
+def _trace_context(
+    metadata: Mapping[object, object],
+) -> dict[str, str]:  # mutable-ok: LiteLLM compatibility payload
+    return {  # mutable-ok: LiteLLM compatibility payload
+        name: str(metadata[name])
+        for name in ("traceparent", "tracestate")
+        if name in metadata and isinstance(metadata[name], str)
+    }
+
+
+def _timestamp(value: object) -> float:
+    if isinstance(value, datetime):
+        return float(value.timestamp())
+    if isinstance(value, int | float):
+        return float(value)
+    return time.time()
+
+
+def _bool_param(params: Mapping[str, object], key: str, default: bool) -> bool:
+    value = params.get(key)  # rebind-ok: invocation-scoped RPC state
+    return value if isinstance(value, bool) else default
+
+
+def _str_param(params: Mapping[str, object], key: str) -> str | None:
+    value = params.get(key)  # rebind-ok: invocation-scoped RPC state
+    return value if isinstance(value, str) else None
+
+
+def _int_param(params: Mapping[str, object], key: str) -> int | None:
+    value = params.get(key)  # rebind-ok: invocation-scoped RPC state
+    return value if isinstance(value, int) and not isinstance(value, bool) else None

litellm/extensions/cache_gateway.py

@@ -0,0 +1,115 @@
+# pyright: reportMissingModuleSource=false, reportMissingTypeArgument=false
+# pyright: reportUnknownMemberType=false, reportUnknownParameterType=false, reportUnknownVariableType=false
+from __future__ import annotations
+
+import json
+import secrets
+from collections.abc import Iterable
+from dataclasses import dataclass
+from typing import Final, Protocol, cast  # noqa: TID251  # validated cache/protobuf boundaries below
+
+import grpc
+
+from litellm.python_extension.generated.v1 import extension_host_pb2 as pb
+from litellm.python_extension.generated.v1 import extension_host_pb2_grpc as pb_grpc
+from litellm.python_extension_host.constants import TOKEN_METADATA_KEY
+
+
+class _AsyncCache(Protocol):
+    async def async_get_cache(
+        self,
+        key: str,
+        local_only: bool = False,
+        **kwargs: object,  # kwargs-ok: LiteLLM callback compatibility
+    ) -> object | None: ...
+
+    async def async_set_cache(
+        self,
+        key: str,
+        value: object,
+        **kwargs: object,  # kwargs-ok: LiteLLM callback compatibility
+    ) -> None: ...
+
+
+@dataclass(frozen=True, slots=True)
+class _CacheBinding:
+    invocation_id: str
+    cache: _AsyncCache
+
+
+class InvocationCacheRegistry:
+    def __init__(self) -> None:
+        self._bindings: dict[str, _CacheBinding] = {}  # mutable-ok: LiteLLM compatibility payload
+
+    def register(self, invocation_id: str, cache: object | None) -> pb.CacheRef | None:
+        if cache is None:
+            return None
+        handle: Final = secrets.token_urlsafe(24)
+        self._bindings[handle] = _CacheBinding(
+            invocation_id,
+            cast(_AsyncCache, cache),  # cast-ok: validated protobuf boundary
+        )
+        return pb.CacheRef(invocation_id=invocation_id, opaque_handle=handle)
+
+    def resolve(self, cache_ref: pb.CacheRef) -> _AsyncCache | None:
+        binding: Final = self._bindings.get(cache_ref.opaque_handle)
+        if binding is None or binding.invocation_id != cache_ref.invocation_id:
+            return None
+        return binding.cache
+
+    def revoke(self, cache_ref: pb.CacheRef | None) -> None:
+        if cache_ref is not None:
+            self._bindings.pop(cache_ref.opaque_handle, None)
+
+
+class GatewayServices(pb_grpc.GatewayServicesServicer):
+    def __init__(self, token: str, registry: InvocationCacheRegistry) -> None:
+        self._token: Final = token
+        self._registry: Final = registry
+
+    async def CacheGet(self, request: pb.CacheGetRequest, context: grpc.aio.ServicerContext) -> pb.CacheGetResponse:
+        await self._authenticate(context)
+        cache = self._registry.resolve(request.cache)  # rebind-ok: invocation-scoped RPC state
+        if cache is None:
+            return pb.CacheGetResponse(operation=_error(pb.ERROR_CODE_NOT_FOUND, "cache reference is invalid"))
+        try:
+            value = await cache.async_get_cache(  # rebind-ok: invocation-scoped RPC state
+                key=request.key, local_only=request.local_only
+            )
+            response: Final = pb.CacheGetResponse(operation=pb.OperationResult(ok=True))
+            if value is not None:
+                response.value_json = json.dumps(value, separators=(",", ":"), default=str).encode()
+        except Exception as error:  # noqa: BLE001  # cache backends may raise arbitrary exceptions
+            return pb.CacheGetResponse(operation=_error(pb.ERROR_CODE_EXTENSION_FAILED, str(error)))
+        else:
+            return response
+
+    async def CacheSet(self, request: pb.CacheSetRequest, context: grpc.aio.ServicerContext) -> pb.OperationResult:
+        await self._authenticate(context)
+        cache = self._registry.resolve(request.cache)  # rebind-ok: invocation-scoped RPC state
+        if cache is None:
+            return _error(pb.ERROR_CODE_NOT_FOUND, "cache reference is invalid")
+        try:
+            kwargs: dict[str, object] = {  # mutable-ok: cache kwargs # rebind-ok: cache kwargs
+                "local_only": request.local_only
+            }
+            if request.HasField("ttl_seconds"):
+                kwargs["ttl"] = request.ttl_seconds
+            value = cast(  # cast-ok: protobuf value # rebind-ok: decoded value
+                object, json.loads(request.value_json)
+            )
+            await cache.async_set_cache(request.key, value, **kwargs)
+            return pb.OperationResult(ok=True)
+        except Exception as error:  # noqa: BLE001  # cache backends may raise arbitrary exceptions
+            return _error(pb.ERROR_CODE_EXTENSION_FAILED, str(error))
+
+    async def _authenticate(self, context: grpc.aio.ServicerContext) -> None:
+        metadata: Final = cast(  # cast-ok: validated protobuf boundary
+            Iterable[tuple[str, str]], context.invocation_metadata() or ()
+        )  # cast-ok: validated protobuf boundary
+        if dict(metadata).get(TOKEN_METADATA_KEY) != self._token:  # mutable-ok: LiteLLM compatibility payload
+            await context.abort(grpc.StatusCode.UNAUTHENTICATED, "invalid extension host token")
+
+
+def _error(code: pb.ErrorCode, message: str) -> pb.OperationResult:
+    return pb.OperationResult(ok=False, error_code=code, error_message=message)

litellm/extensions/client.py

@@ -0,0 +1,283 @@
+# pyright: reportMissingModuleSource=false, reportUnknownMemberType=false, reportUnknownVariableType=false
+from __future__ import annotations
+
+import asyncio
+import contextlib
+from collections.abc import AsyncIterator, Iterable
+from dataclasses import dataclass
+from typing import Final, cast  # noqa: TID251  # protobuf repeated fields lack precise runtime typing
+
+import grpc
+
+from litellm._logging import verbose_proxy_logger
+from litellm.python_extension.generated.v1 import extension_host_pb2 as pb
+from litellm.python_extension.generated.v1 import extension_host_pb2_grpc as pb_grpc
+from litellm.python_extension_host.constants import PROTOCOL_MAJOR, PROTOCOL_MINOR, TOKEN_METADATA_KEY
+
+from .config import ExtensionHostSettings
+from .manifest import ExtensionManifest
+
+
+@dataclass(frozen=True, slots=True)
+class ExtensionHostHealth:
+    healthy: bool
+    reason: str | None = None
+
+
+class PythonExtensionClient:
+    def __init__(self, settings: ExtensionHostSettings, manifest: ExtensionManifest) -> None:
+        self.settings: Final = settings
+        self.manifest: Final = manifest
+        target: Final = settings.endpoint.removeprefix("http://").removeprefix("https://")
+        self._channel: Final = grpc.aio.insecure_channel(target)
+        self._stub: Final = pb_grpc.PythonExtensionHostStub(self._channel)
+        self._metadata: Final = ((TOKEN_METADATA_KEY, settings.token),)
+        self._queue: Final[asyncio.Queue[pb.CallbackEvent]] = asyncio.Queue(settings.callback_queue_size)
+        self._worker: asyncio.Task[None] | None = None
+        self._recovery: asyncio.Task[None] | None = None
+        self._recovery_lock: Final = asyncio.Lock()
+        self._health = ExtensionHostHealth(False, "not connected")
+        self._bypass_counts: dict[tuple[str, str], int] = {}  # mutable-ok: LiteLLM compatibility payload
+        self._descriptor_hooks: dict[str, frozenset[str]] = {}  # mutable-ok: active revision descriptors
+        self._closed = False
+
+    @property
+    def health(self) -> ExtensionHostHealth:
+        return self._health
+
+    @property
+    def bypass_counts(
+        self,
+    ) -> dict[tuple[str, str], int]:  # mutable-ok: LiteLLM compatibility payload
+        return dict(self._bypass_counts)  # mutable-ok: LiteLLM compatibility payload
+
+    def descriptor_hooks(self, plugin_id: str) -> frozenset[str]:
+        return self._descriptor_hooks.get(plugin_id, frozenset())
+
+    async def start(self) -> tuple[pb.ExtensionDescriptor, ...]:
+        descriptors: tuple[pb.ExtensionDescriptor, ...] = ()  # rebind-ok: invocation-scoped RPC state
+        try:
+            await asyncio.wait_for(self._channel.channel_ready(), self.settings.connect_timeout_seconds)
+            descriptors = await self._activate()  # rebind-ok: invocation-scoped RPC state
+        except TimeoutError:
+            self._mark_unhealthy("connect timeout")
+            self._schedule_recovery()
+        except grpc.aio.AioRpcError as error:
+            if error.code() in (grpc.StatusCode.UNAVAILABLE, grpc.StatusCode.DEADLINE_EXCEEDED):
+                self._mark_unhealthy(error.code().name)
+                self._schedule_recovery()
+            else:
+                raise RuntimeError(f"extension host startup failed: {error.details()}") from error
+        self._worker = asyncio.create_task(self._callback_worker(), name="python-extension-callbacks")
+        return descriptors
+
+    async def close(self) -> None:
+        self._closed = True
+        if self._worker is not None:
+            self._worker.cancel()
+            with contextlib.suppress(asyncio.CancelledError):
+                await self._worker
+        if self._recovery is not None:
+            self._recovery.cancel()
+            with contextlib.suppress(asyncio.CancelledError):
+                await self._recovery
+        await self._channel.close()
+
+    async def retire(self, revision_id: str) -> None:
+        try:
+            await self._stub.RetireRevision(
+                pb.RetireRevisionRequest(revision_id=revision_id),
+                timeout=self.settings.hook_timeout_seconds,
+                metadata=self._metadata,
+            )
+        except grpc.aio.AioRpcError:
+            return
+
+    async def execute_guardrail(self, invocation: pb.GuardrailInvocation) -> pb.GuardrailResult:
+        try:
+            result: Final = await self._stub.ExecuteGuardrail(
+                invocation,
+                timeout=self.settings.hook_timeout_seconds,
+                metadata=self._metadata,
+            )
+        except (grpc.aio.AioRpcError, TimeoutError) as error:
+            reason: Final = _rpc_reason(error)
+            self._record_bypass(invocation.plugin_id, pb.HookPhase.Name(invocation.hook_phase), reason)
+            self._schedule_recovery()
+            return pb.GuardrailResult(
+                operation=pb.OperationResult(ok=True),
+                decision=pb.GUARDRAIL_DECISION_ALLOW,
+            )
+        else:
+            if result.operation.ok:
+                self._health = ExtensionHostHealth(True)
+            else:
+                self._record_bypass(
+                    invocation.plugin_id,
+                    pb.HookPhase.Name(invocation.hook_phase),
+                    "extension_error",
+                )
+            return result
+
+    def enqueue_callback(self, event: pb.CallbackEvent) -> bool:
+        try:
+            self._queue.put_nowait(event)
+        except asyncio.QueueFull:
+            self._record_bypass(event.plugin_id, "callback", "queue_full")
+            return False
+        else:
+            return True
+
+    async def transform_stream(self, frames: AsyncIterator[pb.StreamFrame]) -> AsyncIterator[pb.StreamFrame]:
+        call: Final = self._stub.TransformStream(frames, metadata=self._metadata)
+        try:
+            async for frame in call:
+                yield frame
+        except grpc.aio.AioRpcError as error:
+            self._record_bypass("stream", "transform", error.code().name)
+            self._schedule_recovery()
+        finally:
+            if not call.done():
+                call.cancel()
+
+    async def _activate(self) -> tuple[pb.ExtensionDescriptor, ...]:
+        capabilities: Final = await self._stub.GetCapabilities(
+            pb.GetCapabilitiesRequest(protocol_major=PROTOCOL_MAJOR, protocol_minor=PROTOCOL_MINOR),
+            timeout=self.settings.connect_timeout_seconds,
+            metadata=self._metadata,
+        )
+        if capabilities.protocol_major != PROTOCOL_MAJOR:
+            raise RuntimeError(
+                f"extension host protocol major {capabilities.protocol_major} does not match {PROTOCOL_MAJOR}"
+            )
+        has_callbacks: Final = any(spec.kind == pb.EXTENSION_KIND_CALLBACK for spec in self.manifest.specs)
+        if has_callbacks and not capabilities.supports_callback_batching:
+            raise RuntimeError("extension host does not support callback batching")
+        if (
+            capabilities.max_callback_batch_size
+            and self.settings.callback_batch_size > capabilities.max_callback_batch_size
+        ):
+            raise RuntimeError("python_extension_host.callback_batch_size exceeds the host capability")
+        if self.settings.gateway_listen is not None and not capabilities.supports_cache:
+            raise RuntimeError("extension host was not configured for reverse cache access")
+        response: pb.PrepareRevisionResponse = (  # rebind-ok: invocation-scoped RPC state
+            await self._stub.PrepareRevision(  # rebind-ok: invocation-scoped RPC state
+                pb.PrepareRevisionRequest(
+                    revision_id=self.manifest.revision_id,
+                    extensions=self.manifest.specs,
+                ),
+                timeout=self.settings.hook_timeout_seconds,
+                metadata=self._metadata,
+            )
+        )
+        if not response.operation.ok and response.operation.error_code != pb.ERROR_CODE_ALREADY_EXISTS:
+            raise RuntimeError(f"extension manifest rejected: {response.operation.error_message}")
+        streaming_hooks: Final = {  # mutable-ok: LiteLLM compatibility payload
+            "async_post_call_streaming_hook",
+            "async_post_call_streaming_iterator_hook",
+        }
+        if not capabilities.supports_duplex_streaming and any(
+            streaming_hooks.intersection(cast(Iterable[str], descriptor.hooks))  # cast-ok: validated protobuf boundary
+            for descriptor in cast(  # cast-ok: validated protobuf boundary
+                Iterable[pb.ExtensionDescriptor], response.extensions
+            )  # cast-ok: validated protobuf boundary
+        ):
+            raise RuntimeError("extension host does not support required duplex streaming hooks")
+        commit: Final = await self._stub.CommitRevision(
+            pb.CommitRevisionRequest(revision_id=self.manifest.revision_id),
+            timeout=self.settings.hook_timeout_seconds,
+            metadata=self._metadata,
+        )
+        if not commit.ok:
+            raise RuntimeError(f"extension manifest commit failed: {commit.error_message}")
+        self._health = ExtensionHostHealth(True)
+        extensions: Final = cast(  # cast-ok: validated protobuf boundary
+            Iterable[pb.ExtensionDescriptor], response.extensions
+        )  # cast-ok: validated protobuf boundary
+        descriptors: Final = tuple(extensions)
+        self._descriptor_hooks = {  # mutable-ok: active revision descriptors
+            descriptor.id: frozenset(descriptor.hooks) for descriptor in descriptors
+        }
+        return descriptors
+
+    async def _callback_worker(self) -> None:
+        while True:
+            first = await self._queue.get()
+            batch = [first]  # mutable-ok: LiteLLM compatibility payload
+            await asyncio.sleep(0.01)
+            while len(batch) < self.settings.callback_batch_size:
+                try:
+                    batch.append(self._queue.get_nowait())
+                except asyncio.QueueEmpty:
+                    break
+            try:
+                response = await self._stub.PublishCallbackEvents(  # rebind-ok: callback batch response
+                    pb.PublishCallbackEventsRequest(events=batch),
+                    timeout=self.settings.hook_timeout_seconds,
+                    metadata=self._metadata,
+                )
+            except (grpc.aio.AioRpcError, TimeoutError) as error:
+                reason = _rpc_reason(error)
+                for event in batch:
+                    self._record_bypass(event.plugin_id, "callback", reason)
+                self._schedule_recovery()
+            else:
+                operations = tuple(  # cast-ok: validated protobuf repeated field # rebind-ok: callback batch operations
+                    cast(Iterable[pb.OperationResult], response.operations)  # cast-ok: protobuf repeated field
+                )
+                all_ok = len(operations) == len(batch)
+                for index, event in enumerate(batch):
+                    if index >= len(operations):
+                        self._record_bypass(event.plugin_id, "callback", "missing_operation")
+                        continue
+                    operation = operations[index]  # rebind-ok: each callback has one operation
+                    if not operation.ok:
+                        all_ok = False
+                        self._record_bypass(
+                            event.plugin_id,
+                            "callback",
+                            _operation_reason(operation),
+                        )
+                if all_ok:
+                    self._health = ExtensionHostHealth(True)
+            finally:
+                for _ in batch:
+                    self._queue.task_done()
+
+    def _schedule_recovery(self) -> None:
+        if self._closed or (self._recovery is not None and not self._recovery.done()):
+            return
+        self._recovery = asyncio.create_task(self._recover(), name="python-extension-recovery")
+
+    async def _recover(self) -> None:
+        async with self._recovery_lock:
+            delay = 0.25  # rebind-ok: invocation-scoped RPC state
+            while not self._closed:
+                try:
+                    await asyncio.wait_for(self._channel.channel_ready(), self.settings.connect_timeout_seconds)
+                    await self._activate()
+                except (grpc.aio.AioRpcError, RuntimeError, TimeoutError) as error:
+                    self._mark_unhealthy(_rpc_reason(error))
+                    await asyncio.sleep(delay)
+                    delay = min(delay * 2, 5.0)
+                else:
+                    return
+
+    def _record_bypass(self, plugin_id: str, hook: str, reason: str) -> None:
+        key: Final = (plugin_id, f"{hook}:{reason}")
+        self._bypass_counts[key] = self._bypass_counts.get(key, 0) + 1
+        self._mark_unhealthy(reason)
+        verbose_proxy_logger.warning(
+            "python extension host bypass plugin=%s hook=%s reason=%s", plugin_id, hook, reason
+        )
+
+    def _mark_unhealthy(self, reason: str) -> None:
+        self._health = ExtensionHostHealth(False, reason)
+
+
+def _rpc_reason(error: BaseException) -> str:
+    return error.code().name if isinstance(error, grpc.aio.AioRpcError) else str(error)
+
+
+def _operation_reason(operation: pb.OperationResult) -> str:
+    return operation.error_message or "extension_error"

litellm/extensions/config.py

@@ -0,0 +1,56 @@
+# pyright: reportUnknownArgumentType=false, reportUnknownMemberType=false, reportUnknownVariableType=false
+from __future__ import annotations
+
+import os
+from dataclasses import dataclass
+from typing import Final
+
+
+@dataclass(frozen=True, slots=True)
+class ExtensionHostSettings:
+    endpoint: str
+    token: str
+    connect_timeout_seconds: float = 5.0
+    hook_timeout_seconds: float = 30.0
+    callback_queue_size: int = 1_000
+    callback_batch_size: int = 50
+    gateway_listen: str | None = None
+
+
+def settings_from_config(
+    config: dict[str, object],  # mutable-ok: LiteLLM compatibility payload
+) -> ExtensionHostSettings | None:
+    general_settings: Final = config.get("general_settings")
+    if not isinstance(general_settings, dict):
+        return None
+    raw: Final = general_settings.get("python_extension_host")
+    if raw is None:
+        return None
+    if not isinstance(raw, dict):
+        raise TypeError("general_settings.python_extension_host must be an object")
+    endpoint: Final = raw.get("endpoint")
+    token: Final = raw.get("token")
+    if not isinstance(endpoint, str) or not endpoint:
+        raise ValueError("python_extension_host.endpoint is required")
+    if not isinstance(token, str) or not token:
+        raise ValueError("python_extension_host.token is required")
+    resolved_token: Final = _resolve_env(token)
+    return ExtensionHostSettings(
+        endpoint=endpoint,
+        token=resolved_token,
+        connect_timeout_seconds=float(raw.get("connect_timeout_seconds", 5)),
+        hook_timeout_seconds=float(raw.get("hook_timeout_seconds", 30)),
+        callback_queue_size=int(raw.get("callback_queue_size", 1_000)),
+        callback_batch_size=int(raw.get("callback_batch_size", 50)),
+        gateway_listen=(str(raw["gateway_listen"]) if raw.get("gateway_listen") else None),
+    )
+
+
+def _resolve_env(value: str) -> str:
+    if not value.startswith("os.environ/"):
+        return value
+    name: Final = value.removeprefix("os.environ/")
+    resolved: Final = os.environ.get(name)
+    if not resolved:
+        raise ValueError(f"environment variable {name!r} is required for python_extension_host.token")
+    return resolved

litellm/extensions/manifest.py

@@ -0,0 +1,149 @@
+# pyright: reportUnknownArgumentType=false, reportUnknownMemberType=false, reportUnknownVariableType=false
+from __future__ import annotations
+
+import hashlib
+import json
+from collections.abc import Iterable, Mapping
+from dataclasses import dataclass
+from typing import Final
+
+from litellm.python_extension.generated.v1 import extension_host_pb2 as pb
+
+
+@dataclass(frozen=True, slots=True)
+class ExtensionManifest:
+    revision_id: str
+    specs: tuple[pb.ExtensionSpec, ...]
+    callback_ids: Mapping[str, str]
+    guardrail_ids: Mapping[tuple[str, str], str]
+
+
+def build_manifest(config: Mapping[str, object]) -> ExtensionManifest:
+    specs: list[  # mutable-ok: LiteLLM compatibility payload # rebind-ok: invocation-scoped RPC state
+        pb.ExtensionSpec
+    ] = []  # mutable-ok: LiteLLM compatibility payload # rebind-ok: invocation-scoped RPC state
+    callback_ids: dict[  # mutable-ok: LiteLLM compatibility payload # rebind-ok: invocation-scoped RPC state
+        str, str
+    ] = {}  # mutable-ok: LiteLLM compatibility payload # rebind-ok: invocation-scoped RPC state
+    guardrail_ids: dict[  # mutable-ok: LiteLLM compatibility payload # rebind-ok: invocation-scoped RPC state
+        tuple[str, str], str
+    ] = {}  # mutable-ok: LiteLLM compatibility payload # rebind-ok: invocation-scoped RPC state
+    callback_events: dict[  # mutable-ok: LiteLLM compatibility payload # rebind-ok: invocation-scoped RPC state
+        str, set[str]
+    ] = {}  # mutable-ok: LiteLLM compatibility payload # rebind-ok: invocation-scoped RPC state
+    litellm_settings: Final = config.get("litellm_settings")
+    settings: Final = (
+        litellm_settings if isinstance(litellm_settings, Mapping) else {}  # mutable-ok: LiteLLM compatibility payload
+    )  # mutable-ok: LiteLLM compatibility payload
+    for setting, event in (
+        ("callbacks", None),
+        ("success_callback", "success"),
+        ("failure_callback", "failure"),
+    ):
+        for entrypoint in _string_entries(settings.get(setting)):
+            if _is_customer_entrypoint(entrypoint):
+                events = callback_events.setdefault(
+                    entrypoint,
+                    set(),  # mutable-ok: LiteLLM compatibility payload
+                )
+                events.update(("success", "failure") if event is None else (event,))
+    for entrypoint, events in sorted(callback_events.items()):
+        extension_id = _stable_id("callback", entrypoint)
+        callback_ids[entrypoint] = extension_id
+        specs.append(
+            pb.ExtensionSpec(
+                id=extension_id,
+                kind=pb.EXTENSION_KIND_CALLBACK,
+                entrypoint=entrypoint,
+                constructor_json=_canonical_json(
+                    {"callback_events": sorted(events)}  # mutable-ok: LiteLLM compatibility payload
+                ),  # mutable-ok: LiteLLM compatibility payload
+            )
+        )
+    guardrail_configs: Final = _guardrail_configs(config, settings)
+    for guardrail in guardrail_configs:
+        name = guardrail.get("guardrail_name")
+        params = guardrail.get("litellm_params")
+        if not isinstance(name, str) or not isinstance(params, Mapping):
+            continue
+        entrypoint = params.get("guardrail")
+        if not isinstance(entrypoint, str) or not _is_customer_entrypoint(entrypoint):
+            continue
+        extension_id = _stable_id("guardrail", entrypoint, name)
+        guardrail_ids[(entrypoint, name)] = extension_id
+        kwargs = dict(params)  # mutable-ok: LiteLLM compatibility payload
+        kwargs.pop("guardrail", None)
+        mode = kwargs.pop("mode", None)
+        default_on = kwargs.pop("default_on", False)
+        kwargs.update(guardrail_name=name, event_hook=mode, default_on=default_on)
+        specs.append(
+            pb.ExtensionSpec(
+                id=extension_id,
+                kind=pb.EXTENSION_KIND_GUARDRAIL,
+                entrypoint=entrypoint,
+                constructor_json=_canonical_json({"kwargs": kwargs}),  # mutable-ok: LiteLLM compatibility payload
+            )
+        )
+    canonical_specs: Final = b"".join(spec.SerializeToString(deterministic=True) for spec in specs)
+    revision_id: Final = hashlib.sha256(canonical_specs).hexdigest()[:24]
+    return ExtensionManifest(revision_id, tuple(specs), callback_ids, guardrail_ids)
+
+
+def manifest_json_from_config_path(config_path: str) -> str:
+    import yaml
+
+    with open(config_path, encoding="utf-8") as config_file:
+        raw: Final = yaml.safe_load(config_file) or {}  # mutable-ok: LiteLLM compatibility payload
+    if not isinstance(raw, dict):
+        raise TypeError("LiteLLM config must contain an object")
+    manifest: Final = build_manifest(raw)
+    extensions: Final = [  # mutable-ok: LiteLLM compatibility payload
+        {  # mutable-ok: LiteLLM compatibility payload
+            "id": spec.id,
+            "kind": "guardrail" if spec.kind == pb.EXTENSION_KIND_GUARDRAIL else "callback",
+            "entrypoint": spec.entrypoint,
+            "constructor": json.loads(spec.constructor_json or b"{}"),
+        }
+        for spec in manifest.specs
+    ]
+    return json.dumps(
+        {  # mutable-ok: LiteLLM compatibility payload
+            "revision_id": manifest.revision_id,
+            "extensions": extensions,
+        },
+        separators=(",", ":"),
+        sort_keys=True,
+    )
+
+
+def _guardrail_configs(
+    config: Mapping[str, object], settings: Mapping[str, object]
+) -> tuple[Mapping[str, object], ...]:
+    values: list[  # mutable-ok: LiteLLM compatibility payload # rebind-ok: invocation-scoped RPC state
+        Mapping[str, object]
+    ] = []  # mutable-ok: LiteLLM compatibility payload # rebind-ok: invocation-scoped RPC state
+    for candidate in (config.get("guardrails"), settings.get("guardrails")):
+        if isinstance(candidate, list):
+            values.extend(item for item in candidate if isinstance(item, Mapping))
+    return tuple(values)
+
+
+def _string_entries(value: object) -> Iterable[str]:
+    if isinstance(value, str):
+        return (value,)
+    if isinstance(value, list):
+        return (item for item in value if isinstance(item, str))
+    return ()
+
+
+def _is_customer_entrypoint(value: str) -> bool:
+    return ("." in value or ":" in value) and not value.startswith(("http://", "https://"))
+
+
+def _stable_id(kind: str, entrypoint: str, name: str = "") -> str:
+    digest: Final = hashlib.sha256(f"{kind}\0{entrypoint}\0{name}".encode()).hexdigest()[:16]
+    return f"{kind}-{digest}"
+
+
+def _canonical_json(value: object) -> bytes:
+    return json.dumps(value, separators=(",", ":"), sort_keys=True, default=str).encode()

litellm/extensions/runtime.py

@@ -0,0 +1,138 @@
+# pyright: reportMissingModuleSource=false, reportUnknownMemberType=false
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Final
+
+import grpc
+
+from litellm.python_extension.generated.v1 import extension_host_pb2_grpc as pb_grpc
+from litellm.types.guardrails import GuardrailEventHooks, Mode
+
+from .adapters import RemoteCustomGuardrail, RemoteCustomLogger
+from .cache_gateway import GatewayServices, InvocationCacheRegistry
+from .client import PythonExtensionClient
+from .config import ExtensionHostSettings, settings_from_config
+from .manifest import ExtensionManifest, build_manifest
+
+
+@dataclass(slots=True)
+class ExtensionRuntime:
+    settings: ExtensionHostSettings
+    manifest: ExtensionManifest
+    client: PythonExtensionClient
+    cache_registry: InvocationCacheRegistry
+    gateway_server: grpc.aio.Server | None = None
+
+    def callback(self, entrypoint: str) -> RemoteCustomLogger | None:
+        plugin_id = self.manifest.callback_ids.get(entrypoint)  # rebind-ok: invocation-scoped RPC state
+        if plugin_id is None:
+            return None
+        return RemoteCustomLogger(self.client, plugin_id, self.cache_registry)
+
+    def guardrail(
+        self,
+        entrypoint: str,
+        name: str,
+        event_hook: GuardrailEventHooks  # mutable-ok: LiteLLM compatibility payload
+        | list[GuardrailEventHooks]
+        | Mode
+        | None = None,  # mutable-ok: LiteLLM compatibility payload
+        default_on: bool = False,
+        **kwargs: object,  # kwargs-ok: LiteLLM callback compatibility
+    ) -> RemoteCustomGuardrail | None:
+        plugin_id = self.manifest.guardrail_ids.get((entrypoint, name))  # rebind-ok: invocation-scoped RPC state
+        if plugin_id is None:
+            return None
+        return RemoteCustomGuardrail(
+            self.client,
+            plugin_id,
+            self.cache_registry,
+            guardrail_name=name,
+            event_hook=event_hook,
+            default_on=default_on,
+            extra_params=kwargs,
+        )
+
+    async def close(self) -> None:
+        await self.client.close()
+        if self.gateway_server is not None:
+            await self.gateway_server.stop(grace=5)
+
+
+_runtime: ExtensionRuntime | None = None  # rebind-ok: invocation-scoped RPC state
+
+
+def get_extension_runtime() -> ExtensionRuntime | None:
+    return _runtime
+
+
+async def configure_extension_runtime(
+    config: dict[str, object],  # mutable-ok: LiteLLM compatibility payload
+) -> ExtensionRuntime | None:
+    global _runtime  # noqa: PLW0603  # process-wide extension lifecycle singleton
+    settings: Final = settings_from_config(config)
+    if settings is None:
+        if _runtime is not None:
+            await _runtime.close()
+            _runtime = None  # rebind-ok: invocation-scoped RPC state
+        return None
+    manifest: Final = build_manifest(config)
+    cache_registry: Final = InvocationCacheRegistry()
+    gateway_server: Final = await _start_gateway_services(settings, cache_registry)
+    client: Final = PythonExtensionClient(settings, manifest)
+    try:
+        await client.start()
+    except Exception:
+        await client.close()
+        if gateway_server is not None:
+            await gateway_server.stop(grace=0)
+        raise
+    previous: Final = _runtime
+    _runtime = ExtensionRuntime(  # rebind-ok: invocation-scoped RPC state
+        settings, manifest, client, cache_registry, gateway_server
+    )  # rebind-ok: invocation-scoped RPC state
+    if previous is not None:
+        if previous.manifest.revision_id != manifest.revision_id:
+            await previous.client.retire(previous.manifest.revision_id)
+        await previous.close()
+    return _runtime
+
+
+def remote_callback(entrypoint: str) -> RemoteCustomLogger | None:
+    return _runtime.callback(entrypoint) if _runtime is not None else None
+
+
+def remote_guardrail(
+    entrypoint: str,
+    name: str,
+    event_hook: GuardrailEventHooks  # mutable-ok: LiteLLM compatibility payload
+    | list[GuardrailEventHooks]
+    | Mode
+    | None = None,  # mutable-ok: LiteLLM compatibility payload
+    default_on: bool = False,
+    **kwargs: object,  # kwargs-ok: LiteLLM callback compatibility
+) -> RemoteCustomGuardrail | None:
+    if _runtime is None:
+        return None
+    return _runtime.guardrail(
+        entrypoint,
+        name,
+        event_hook=event_hook,
+        default_on=default_on,
+        **kwargs,
+    )
+
+
+async def _start_gateway_services(
+    settings: ExtensionHostSettings, registry: InvocationCacheRegistry
+) -> grpc.aio.Server | None:
+    if settings.gateway_listen is None:
+        return None
+    server: Final = grpc.aio.server()
+    pb_grpc.add_GatewayServicesServicer_to_server(GatewayServices(settings.token, registry), server)
+    target: Final = settings.gateway_listen.removeprefix("http://").removeprefix("https://")
+    if server.add_insecure_port(target) == 0:
+        raise RuntimeError(f"failed to bind GatewayServices to {settings.gateway_listen}")
+    await server.start()
+    return server

litellm/proxy/common_utils/callback_utils.py

@@ -369,15 +369,25 @@ def initialize_callbacks_on_proxy(
                 verbose_proxy_logger.debug(
                     "%s attempting to import custom calback=%s %s", blue_color_code, callback, reset_color_code
                 )
-                imported_list.append(
-                    _loaded_callback_or_raise(
-                        entry=callback,
-                        loaded=get_instance_fn(
-                            value=callback,
-                            config_file_path=config_file_path,
-                        ),
-                    )
+                from litellm.extensions.runtime import get_extension_runtime, remote_callback
+
+                remote_for_callback = remote_callback(  # rebind-ok: each configured callback is independent
+                    callback
                 )
+                if remote_for_callback is not None:
+                    imported_list.append(remote_for_callback)
+                elif get_extension_runtime() is not None:
+                    raise ValueError(f"callback {callback!r} was not present in the extension host manifest")
+                else:
+                    imported_list.append(
+                        _loaded_callback_or_raise(
+                            entry=callback,
+                            loaded=get_instance_fn(
+                                value=callback,
+                                config_file_path=config_file_path,
+                            ),
+                        )
+                    )
         if isinstance(litellm.callbacks, list):
             litellm.callbacks.extend(imported_list)
         else:
@@ -388,15 +398,23 @@ def initialize_callbacks_on_proxy(
 
             PrometheusLogger._mount_metrics_endpoint()
     else:
-        litellm.callbacks = [
-            _loaded_callback_or_raise(
-                entry=value,
-                loaded=get_instance_fn(
-                    value=value,
-                    config_file_path=config_file_path,
-                ),
-            )
-        ]
+        from litellm.extensions.runtime import get_extension_runtime, remote_callback
+
+        remote_single: Final = remote_callback(value)
+        if remote_single is not None:
+            litellm.callbacks = [remote_single]  # mutable-ok: global callback registry requires a list
+        elif get_extension_runtime() is not None:
+            raise ValueError(f"callback {value!r} was not present in the extension host manifest")
+        else:
+            litellm.callbacks = [  # mutable-ok: global callback registry requires a list
+                _loaded_callback_or_raise(
+                    entry=value,
+                    loaded=get_instance_fn(
+                        value=value,
+                        config_file_path=config_file_path,
+                    ),
+                )
+            ]
     verbose_proxy_logger.debug("%s Initialized Callbacks - %s %s", blue_color_code, litellm.callbacks, reset_color_code)
 
 

litellm/proxy/guardrails/guardrail_registry.py

@@ -47,6 +47,7 @@ from litellm.types.guardrails import (
     GuardrailEventHooks,
     LakeraCategoryThresholds,
     LitellmParams,
+    Mode,
     SupportedGuardrailIntegrations,
 )
 
@@ -424,6 +425,16 @@ def _apply_configured_bool_overrides(instance: CustomGuardrail, litellm_params:
         instance.scan_raw_request = bool(litellm_params.scan_raw_request)
 
 
+def _remote_guardrail_event_hook(
+    mode: str | list[str] | Mode,  # mutable-ok: guardrail config accepts a hook list
+) -> GuardrailEventHooks | list[GuardrailEventHooks] | Mode:  # mutable-ok: CustomGuardrail requires a list
+    if isinstance(mode, Mode):
+        return mode
+    if isinstance(mode, list):
+        return [GuardrailEventHooks(hook) for hook in mode]  # mutable-ok: CustomGuardrail requires a list
+    return GuardrailEventHooks(mode)
+
+
 class InMemoryGuardrailHandler:
     """
     Class that handles initializing guardrails and adding them to the CallbackManager
@@ -581,10 +592,6 @@ class InMemoryGuardrailHandler:
             guardrail_type,
         )
 
-        _guardrail_class: Final[Callable[..., CustomGuardrail]] = get_instance_fn(
-            guardrail_type, config_file_path=config_file_path
-        )
-
         mode: Final = litellm_params.mode
         if mode is None:
             raise ValueError(
@@ -605,6 +612,25 @@ class InMemoryGuardrailHandler:
         for key in ["guardrail", "mode", "default_on"]:
             extra_params.pop(key, None)
 
+        from litellm.extensions.runtime import get_extension_runtime, remote_guardrail
+
+        remote: Final = remote_guardrail(
+            guardrail_type,
+            guardrail["guardrail_name"],
+            event_hook=_remote_guardrail_event_hook(mode),
+            default_on=bool(default_on),
+            **extra_params,
+        )
+        if remote is not None:
+            litellm.logging_callback_manager.add_litellm_callback(remote)
+            return remote
+        if get_extension_runtime() is not None:
+            raise ValueError(f"guardrail {guardrail_type!r} was not present in the extension host manifest")
+
+        _guardrail_class: Final[Callable[..., CustomGuardrail]] = get_instance_fn(
+            guardrail_type, config_file_path=config_file_path
+        )
+
         _guardrail_callback: Final = _guardrail_class(
             guardrail_name=guardrail["guardrail_name"],
             event_hook=mode,

litellm/proxy/proxy_server.py

@@ -5044,6 +5044,10 @@ class ProxyConfig:
 
         self._load_environment_variables(config=config)
 
+        from litellm.extensions.runtime import configure_extension_runtime
+
+        await configure_extension_runtime(config)
+
         ## Coordination Redis (before cache init, so the explicit block wins)
         coordination_redis_cache: Final = self._init_coordination_redis(config=config)
         if coordination_redis_cache is not None:
@@ -5246,8 +5250,16 @@ class ProxyConfig:
                     for callback in value:
                         # user passed custom_callbacks.async_on_succes_logger. They need us to import a function
                         if "." in callback:
+                            from litellm.extensions.runtime import get_extension_runtime, remote_callback
+
+                            callback_instance = remote_callback(callback)
+                            if callback_instance is None and get_extension_runtime() is not None:
+                                raise ValueError(
+                                    f"callback {callback!r} was not present in the extension host manifest"
+                                )
                             litellm.logging_callback_manager.add_litellm_success_callback(
-                                get_instance_fn(
+                                callback_instance
+                                or get_instance_fn(
                                     value=callback,
                                     config_file_path=config_file_path,
                                 )
@@ -5273,8 +5285,16 @@ class ProxyConfig:
                     for callback in value:
                         # user passed custom_callbacks.async_on_succes_logger. They need us to import a function
                         if "." in callback:
+                            from litellm.extensions.runtime import get_extension_runtime, remote_callback
+
+                            callback_instance = remote_callback(callback)
+                            if callback_instance is None and get_extension_runtime() is not None:
+                                raise ValueError(
+                                    f"callback {callback!r} was not present in the extension host manifest"
+                                )
                             litellm.logging_callback_manager.add_litellm_failure_callback(
-                                get_instance_fn(
+                                callback_instance
+                                or get_instance_fn(
                                     value=callback,
                                     config_file_path=config_file_path,
                                 )

tests/test_litellm/python_extension_host/__init__.py

@@ -0,0 +1 @@
+"""Python extension host tests."""

tests/test_litellm/python_extension_host/test_rpc_sidecar.py

@@ -0,0 +1,537 @@
+from __future__ import annotations
+
+import asyncio
+import json
+import os
+import socket
+import sys
+from collections.abc import AsyncIterator
+from datetime import datetime
+from pathlib import Path
+
+import grpc
+import pytest
+import pytest_asyncio
+
+from litellm.exceptions import GuardrailRaisedException
+from litellm.extensions.cache_gateway import GatewayServices, InvocationCacheRegistry
+from litellm.extensions.client import PythonExtensionClient
+from litellm.extensions.config import ExtensionHostSettings, settings_from_config
+from litellm.extensions.manifest import build_manifest
+from litellm.extensions.runtime import configure_extension_runtime
+from litellm.integrations.custom_guardrail import CustomGuardrail
+from litellm.integrations.custom_logger import CustomLogger
+from litellm.python_extension.generated.v1 import extension_host_pb2 as pb
+from litellm.python_extension.generated.v1 import extension_host_pb2_grpc as pb_grpc
+from litellm.python_extension_host.constants import TOKEN_METADATA_KEY
+from litellm.python_extension_host.service import PythonExtensionHostService
+
+MODULE = "tests.test_litellm.python_extension_host.test_rpc_sidecar"
+TOKEN = "test-extension-token"
+METADATA = ((TOKEN_METADATA_KEY, TOKEN),)
+
+
+class GuardrailFixture(CustomGuardrail):
+    async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type):
+        data["hosted"] = True
+        data["received_plaintext_key"] = bool(user_api_key_dict.api_key)
+        return data
+
+    async def async_post_call_success_hook(self, data, user_api_key_dict, response):
+        response["post_call"] = data["model"]
+        return response
+
+
+class BlockingGuardrailFixture(CustomGuardrail):
+    async def async_moderation_hook(self, data, user_api_key_dict, call_type):
+        raise GuardrailRaisedException(
+            guardrail_name="blocking",
+            message="blocked by fixture",
+            blocked_content=True,
+        )
+
+
+class CacheGuardrailFixture(CustomGuardrail):
+    async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type):
+        data["cached"] = await cache.async_get_cache("seed")
+        await cache.async_set_cache("written", {"ok": True}, ttl=30)
+        return data
+
+
+class FakeCache:
+    def __init__(self):
+        self.values: dict[str, object] = {"seed": {"value": 7}}
+
+    async def async_get_cache(self, key: str, **kwargs: object) -> object | None:
+        return self.values.get(key)
+
+    async def async_set_cache(self, key: str, value: object, **kwargs: object) -> None:
+        self.values[key] = value
+
+
+class UnsupportedFixture(CustomLogger):
+    async def async_pre_request_hook(self, model, messages, kwargs):
+        return kwargs
+
+
+class CallbackFixture(CustomLogger):
+    events: list[tuple[str, str]] = []
+
+    async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
+        self.events.append((kwargs["model"], response_obj["id"]))
+
+
+function_events: list[str] = []
+
+
+async def callback_function(kwargs, response_obj, start_time, end_time):
+    function_events.append(kwargs["model"])
+
+
+class StreamFixture(CustomLogger):
+    async def async_post_call_streaming_iterator_hook(
+        self, user_api_key_dict, response, request_data
+    ) -> AsyncIterator[dict[str, str]]:
+        async for chunk in response:
+            yield {"value": chunk["value"].upper()}
+            yield {"value": "!"}
+
+
+class ChunkStreamFixture(CustomLogger):
+    async def async_post_call_streaming_hook(self, user_api_key_dict, response):
+        return {"value": response["value"].upper()}
+
+
+@pytest_asyncio.fixture(loop_scope="function")
+async def host_stub():
+    server = grpc.aio.server()
+    pb_grpc.add_PythonExtensionHostServicer_to_server(PythonExtensionHostService(TOKEN), server)
+    port = server.add_insecure_port("127.0.0.1:0")
+    await server.start()
+    channel = grpc.aio.insecure_channel(f"127.0.0.1:{port}")
+    await channel.channel_ready()
+    try:
+        yield pb_grpc.PythonExtensionHostStub(channel)
+    finally:
+        await channel.close()
+        await server.stop(grace=0)
+
+
+@pytest.mark.asyncio
+async def test_capabilities_require_auth_and_negotiate_version(host_stub):
+    with pytest.raises(grpc.aio.AioRpcError) as error:
+        await host_stub.GetCapabilities(pb.GetCapabilitiesRequest(protocol_major=1, protocol_minor=0))
+    assert error.value.code() == grpc.StatusCode.UNAUTHENTICATED
+
+    capabilities = await host_stub.GetCapabilities(
+        pb.GetCapabilitiesRequest(protocol_major=1, protocol_minor=3), metadata=METADATA
+    )
+    assert capabilities.protocol_major == 1
+    assert capabilities.protocol_minor == 0
+    assert capabilities.supports_duplex_streaming is True
+    assert "async_pre_call_hook" in capabilities.supported_hooks
+
+
+@pytest.mark.asyncio
+async def test_prepare_commit_execute_replace_and_retire(host_stub):
+    prepare = await host_stub.PrepareRevision(
+        pb.PrepareRevisionRequest(
+            revision_id="revision-1",
+            extensions=[_spec("guardrail", pb.EXTENSION_KIND_GUARDRAIL, "GuardrailFixture")],
+        ),
+        metadata=METADATA,
+    )
+    assert prepare.operation.ok is True
+    assert prepare.extensions[0].hooks == ["async_post_call_success_hook", "async_pre_call_hook"]
+
+    prepared_again = await host_stub.PrepareRevision(
+        pb.PrepareRevisionRequest(
+            revision_id="revision-1",
+            extensions=[_spec("guardrail", pb.EXTENSION_KIND_GUARDRAIL, "GuardrailFixture")],
+        ),
+        metadata=METADATA,
+    )
+    assert prepared_again.operation.error_code == pb.ERROR_CODE_ALREADY_EXISTS
+    assert prepared_again.extensions[0].hooks == prepare.extensions[0].hooks
+
+    inactive = await host_stub.ExecuteGuardrail(_guardrail_request("revision-1"), metadata=METADATA)
+    assert inactive.decision == pb.GUARDRAIL_DECISION_ERROR
+
+    committed = await host_stub.CommitRevision(pb.CommitRevisionRequest(revision_id="revision-1"), metadata=METADATA)
+    assert committed.ok is True
+    replaced = await host_stub.ExecuteGuardrail(_guardrail_request("revision-1"), metadata=METADATA)
+    assert replaced.decision == pb.GUARDRAIL_DECISION_REPLACE_REQUEST
+    replaced_request = json.loads(replaced.request_json)
+    assert replaced_request["hosted"] is True
+    assert replaced_request["received_plaintext_key"] is False
+
+    active_retire = await host_stub.RetireRevision(
+        pb.RetireRevisionRequest(revision_id="revision-1"), metadata=METADATA
+    )
+    assert active_retire.error_code == pb.ERROR_CODE_INVALID_ARGUMENT
+
+    await host_stub.PrepareRevision(
+        pb.PrepareRevisionRequest(revision_id="revision-2", extensions=[]), metadata=METADATA
+    )
+    await host_stub.CommitRevision(pb.CommitRevisionRequest(revision_id="revision-2"), metadata=METADATA)
+    retired = await host_stub.RetireRevision(pb.RetireRevisionRequest(revision_id="revision-1"), metadata=METADATA)
+    assert retired.ok is True
+
+
+@pytest.mark.asyncio
+async def test_recognized_guardrail_exception_becomes_block(host_stub):
+    await _activate(
+        host_stub,
+        "block-revision",
+        [_spec("blocking", pb.EXTENSION_KIND_GUARDRAIL, "BlockingGuardrailFixture")],
+    )
+    request = _guardrail_request("block-revision", plugin_id="blocking")
+    request.hook_phase = pb.HOOK_PHASE_DURING_CALL
+    result = await host_stub.ExecuteGuardrail(request, metadata=METADATA)
+    assert result.operation.ok is True
+    assert result.decision == pb.GUARDRAIL_DECISION_BLOCK
+    assert result.public_error.status_code == 400
+    assert "blocked by fixture" in result.public_error.message
+
+
+@pytest.mark.asyncio
+async def test_prepare_rejects_any_unsupported_override(host_stub):
+    response = await host_stub.PrepareRevision(
+        pb.PrepareRevisionRequest(
+            revision_id="unsupported",
+            extensions=[_spec("unsupported", pb.EXTENSION_KIND_CALLBACK, "UnsupportedFixture")],
+        ),
+        metadata=METADATA,
+    )
+    assert response.operation.ok is False
+    assert response.operation.error_code == pb.ERROR_CODE_LOAD_FAILED
+    assert "async_pre_request_hook" in response.operation.error_message
+
+
+@pytest.mark.asyncio
+async def test_callback_batch_supports_logger_and_function(host_stub):
+    CallbackFixture.events.clear()
+    function_events.clear()
+    specs = [
+        _spec("logger", pb.EXTENSION_KIND_CALLBACK, "CallbackFixture"),
+        _spec(
+            "function",
+            pb.EXTENSION_KIND_CALLBACK,
+            "callback_function",
+            {"callback_events": ["success"]},
+        ),
+    ]
+    await _activate(host_stub, "callbacks", specs)
+    events = [_callback_event("callbacks", plugin_id) for plugin_id in ("logger", "function")]
+    response = await host_stub.PublishCallbackEvents(pb.PublishCallbackEventsRequest(events=events), metadata=METADATA)
+    assert all(operation.ok for operation in response.operations)
+    assert CallbackFixture.events == [("test-model", "response-1")]
+    assert function_events == ["test-model"]
+
+
+@pytest.mark.asyncio
+async def test_duplex_iterator_can_emit_multiple_chunks_per_input(host_stub):
+    await _activate(
+        host_stub,
+        "streaming",
+        [_spec("stream", pb.EXTENSION_KIND_CALLBACK, "StreamFixture")],
+    )
+
+    async def frames() -> AsyncIterator[pb.StreamFrame]:
+        yield pb.StreamFrame(
+            kind=pb.STREAM_FRAME_KIND_OPEN,
+            stream_id="stream-1",
+            open=pb.StreamOpen(
+                context=_context("streaming"),
+                plugin_id="stream",
+                request_json=b'{"model":"test-model"}',
+                auth=pb.AuthContext(key_hash="hashed"),
+                iterator_hook=True,
+            ),
+        )
+        yield pb.StreamFrame(
+            kind=pb.STREAM_FRAME_KIND_INPUT_CHUNK,
+            stream_id="stream-1",
+            chunk_json=b'{"value":"hello"}',
+        )
+        yield pb.StreamFrame(kind=pb.STREAM_FRAME_KIND_END, stream_id="stream-1")
+
+    output = [frame async for frame in host_stub.TransformStream(frames(), metadata=METADATA)]
+    assert [json.loads(frame.chunk_json) for frame in output[:-1]] == [
+        {"value": "HELLO"},
+        {"value": "!"},
+    ], [(frame.kind, frame.error.message) for frame in output]
+    assert output[-1].kind == pb.STREAM_FRAME_KIND_END
+
+
+@pytest.mark.asyncio
+async def test_duplex_chunk_hook_transforms_each_input(host_stub):
+    await _activate(
+        host_stub,
+        "chunk-streaming",
+        [_spec("chunk-stream", pb.EXTENSION_KIND_CALLBACK, "ChunkStreamFixture")],
+    )
+
+    async def frames() -> AsyncIterator[pb.StreamFrame]:
+        yield pb.StreamFrame(
+            kind=pb.STREAM_FRAME_KIND_OPEN,
+            stream_id="stream-2",
+            open=pb.StreamOpen(
+                context=_context("chunk-streaming"),
+                plugin_id="chunk-stream",
+                request_json=b'{"model":"test-model"}',
+                auth=pb.AuthContext(key_hash="hashed"),
+                iterator_hook=False,
+            ),
+        )
+        yield pb.StreamFrame(
+            kind=pb.STREAM_FRAME_KIND_INPUT_CHUNK,
+            stream_id="stream-2",
+            chunk_json=b'{"value":"hello"}',
+        )
+        yield pb.StreamFrame(kind=pb.STREAM_FRAME_KIND_END, stream_id="stream-2")
+
+    output = [frame async for frame in host_stub.TransformStream(frames(), metadata=METADATA)]
+    assert json.loads(output[0].chunk_json) == {"value": "HELLO"}
+    assert output[1].kind == pb.STREAM_FRAME_KIND_END
+
+
+@pytest.mark.asyncio
+async def test_invocation_scoped_reverse_cache_access_is_revoked():
+    registry = InvocationCacheRegistry()
+    cache = FakeCache()
+    cache_ref = registry.register("invocation-1", cache)
+    assert cache_ref is not None
+
+    gateway_server = grpc.aio.server()
+    pb_grpc.add_GatewayServicesServicer_to_server(GatewayServices(TOKEN, registry), gateway_server)
+    gateway_port = gateway_server.add_insecure_port("127.0.0.1:0")
+    await gateway_server.start()
+    gateway_channel = grpc.aio.insecure_channel(f"127.0.0.1:{gateway_port}")
+    gateway_stub = pb_grpc.GatewayServicesStub(gateway_channel)
+
+    host_server = grpc.aio.server()
+    pb_grpc.add_PythonExtensionHostServicer_to_server(
+        PythonExtensionHostService(TOKEN, gateway_stub=gateway_stub), host_server
+    )
+    host_port = host_server.add_insecure_port("127.0.0.1:0")
+    await host_server.start()
+    host_channel = grpc.aio.insecure_channel(f"127.0.0.1:{host_port}")
+    host = pb_grpc.PythonExtensionHostStub(host_channel)
+    try:
+        await _activate(
+            host,
+            "cache-revision",
+            [_spec("cache", pb.EXTENSION_KIND_GUARDRAIL, "CacheGuardrailFixture")],
+        )
+        request = _guardrail_request("cache-revision", plugin_id="cache")
+        request.cache.CopyFrom(cache_ref)
+        result = await host.ExecuteGuardrail(request, metadata=METADATA)
+        assert json.loads(result.request_json)["cached"] == {"value": 7}
+        assert cache.values["written"] == {"ok": True}
+
+        registry.revoke(cache_ref)
+        revoked = await gateway_stub.CacheGet(pb.CacheGetRequest(cache=cache_ref, key="seed"), metadata=METADATA)
+        assert revoked.operation.error_code == pb.ERROR_CODE_NOT_FOUND
+    finally:
+        await host_channel.close()
+        await host_server.stop(grace=0)
+        await gateway_channel.close()
+        await gateway_server.stop(grace=0)
+
+
+def test_manifest_build_does_not_import_customer_modules():
+    module_name = "customer_extension_that_proxy_must_not_import"
+    config = {
+        "litellm_settings": {
+            "callbacks": [f"{module_name}.logger"],
+            "success_callback": [f"{module_name}.success"],
+            "failure_callback": [f"{module_name}.failure"],
+        },
+        "guardrails": [
+            {
+                "guardrail_name": "customer-guardrail",
+                "litellm_params": {
+                    "guardrail": f"{module_name}.Guardrail",
+                    "mode": "pre_call",
+                },
+            }
+        ],
+    }
+
+    manifest = build_manifest(config)
+
+    assert len(manifest.specs) == 4
+    assert module_name not in sys.modules
+    assert settings_from_config(config) is None
+
+
+@pytest.mark.asyncio
+async def test_python_client_fails_open_when_host_is_unavailable():
+    port = _unused_port()
+    manifest = build_manifest({"litellm_settings": {"callbacks": [f"{MODULE}.CallbackFixture"]}})
+    client = PythonExtensionClient(
+        ExtensionHostSettings(
+            endpoint=f"http://127.0.0.1:{port}",
+            token=TOKEN,
+            connect_timeout_seconds=0.05,
+            hook_timeout_seconds=0.05,
+        ),
+        manifest,
+    )
+    try:
+        assert await client.start() == ()
+        result = await client.execute_guardrail(_guardrail_request(manifest.revision_id))
+        assert result.decision == pb.GUARDRAIL_DECISION_ALLOW
+        assert client.health.healthy is False
+        assert client.bypass_counts
+    finally:
+        await client.close()
+
+
+@pytest.mark.asyncio
+async def test_proxy_and_sidecar_share_contract_without_proxy_import(tmp_path: Path):
+    module_name = "customer_sidecar_only_plugin"
+    import_marker = tmp_path / "import.pid"
+    event_marker = tmp_path / "event.pid"
+    (tmp_path / f"{module_name}.py").write_text(
+        """
+import os
+from pathlib import Path
+from litellm.integrations.custom_logger import CustomLogger
+
+Path(os.environ["EXTENSION_IMPORT_MARKER"]).write_text(str(os.getpid()))
+
+class HostedLogger(CustomLogger):
+    async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
+        Path(os.environ["EXTENSION_EVENT_MARKER"]).write_text(str(os.getpid()))
+""".lstrip(),
+        encoding="utf-8",
+    )
+    port = _unused_port()
+    environment = os.environ.copy()
+    environment["PYTHONPATH"] = os.pathsep.join((str(tmp_path), os.getcwd()))
+    environment["LITELLM_EXTENSION_HOST_TOKEN"] = TOKEN
+    environment["EXTENSION_IMPORT_MARKER"] = str(import_marker)
+    environment["EXTENSION_EVENT_MARKER"] = str(event_marker)
+    process = await asyncio.create_subprocess_exec(
+        sys.executable,
+        "-m",
+        "litellm.python_extension_host.server",
+        "--listen",
+        f"127.0.0.1:{port}",
+        env=environment,
+        stdout=asyncio.subprocess.DEVNULL,
+        stderr=asyncio.subprocess.DEVNULL,
+    )
+    channel = grpc.aio.insecure_channel(f"127.0.0.1:{port}")
+    try:
+        await asyncio.wait_for(channel.channel_ready(), timeout=10)
+        config = {
+            "general_settings": {
+                "python_extension_host": {
+                    "endpoint": f"http://127.0.0.1:{port}",
+                    "token": TOKEN,
+                    "connect_timeout_seconds": 2,
+                    "hook_timeout_seconds": 2,
+                }
+            },
+            "litellm_settings": {
+                "success_callback": [f"{module_name}.HostedLogger"],
+            },
+        }
+        runtime = await configure_extension_runtime(config)
+        assert runtime is not None
+        assert module_name not in sys.modules
+        await _wait_for_file(import_marker)
+        assert int(import_marker.read_text()) == process.pid
+
+        callback = runtime.callback(f"{module_name}.HostedLogger")
+        assert callback is not None
+        await callback.async_log_success_event(
+            {"model": "hosted-model"},
+            {"id": "response-1"},
+            datetime.now(),
+            datetime.now(),
+        )
+        await _wait_for_file(event_marker)
+        assert int(event_marker.read_text()) == process.pid
+        assert process.pid != os.getpid()
+    finally:
+        await configure_extension_runtime({})
+        await channel.close()
+        process.terminate()
+        try:
+            await asyncio.wait_for(process.wait(), timeout=5)
+        except TimeoutError:
+            process.kill()
+            await process.wait()
+
+
+def _unused_port() -> int:
+    with socket.socket() as server_socket:
+        server_socket.bind(("127.0.0.1", 0))
+        return int(server_socket.getsockname()[1])
+
+
+async def _wait_for_file(path: Path) -> None:
+    for _ in range(200):
+        if path.exists():
+            return
+        await asyncio.sleep(0.025)
+    raise AssertionError(f"timed out waiting for {path}")
+
+
+def _spec(
+    plugin_id: str,
+    kind: int,
+    target: str,
+    constructor: dict[str, object] | None = None,
+) -> pb.ExtensionSpec:
+    return pb.ExtensionSpec(
+        id=plugin_id,
+        kind=kind,
+        entrypoint=f"{MODULE}.{target}",
+        constructor_json=json.dumps(constructor or {}).encode(),
+    )
+
+
+async def _activate(host_stub, revision: str, specs: list[pb.ExtensionSpec]) -> None:
+    prepared = await host_stub.PrepareRevision(
+        pb.PrepareRevisionRequest(revision_id=revision, extensions=specs), metadata=METADATA
+    )
+    assert prepared.operation.ok, prepared.operation.error_message
+    committed = await host_stub.CommitRevision(pb.CommitRevisionRequest(revision_id=revision), metadata=METADATA)
+    assert committed.ok
+
+
+def _context(revision: str) -> pb.InvocationContext:
+    return pb.InvocationContext(
+        request_id="request-1",
+        invocation_id="invocation-1",
+        active_revision=revision,
+        api_surface="chat.completions",
+        call_type="acompletion",
+    )
+
+
+def _guardrail_request(revision: str, plugin_id: str = "guardrail") -> pb.GuardrailInvocation:
+    return pb.GuardrailInvocation(
+        context=_context(revision),
+        plugin_id=plugin_id,
+        hook_phase=pb.HOOK_PHASE_PRE_CALL,
+        request_json=b'{"model":"test-model"}',
+        auth=pb.AuthContext(key_hash="sha256-only", user_id="user-1", team_id="team-1"),
+    )
+
+
+def _callback_event(revision: str, plugin_id: str) -> pb.CallbackEvent:
+    return pb.CallbackEvent(
+        context=_context(revision),
+        plugin_id=plugin_id,
+        kind=pb.CALLBACK_EVENT_KIND_SUCCESS,
+        standard_logging_payload_json=b'{"model":"test-model"}',
+        response_json=b'{"id":"response-1"}',
+        start_time_seconds=datetime(2026, 1, 1).timestamp(),
+        end_time_seconds=datetime(2026, 1, 1, 0, 0, 1).timestamp(),
+    )