feat: add proxy function support

pyproject.toml

@@ -29,6 +29,7 @@ dev = [
   "pytest-asyncio>=1.1.0",
   "pytest-cov>=6.2.1",
   "pytest-env>=1.1.5",
+  "pytest-order>=1.3.0",
   "pytest-recording>=0.13.4",
   "ruff==0.14.9",
   "types-pytz>=2025.2.0.20250516",

pytest.ini

@@ -26,10 +26,13 @@ env =
     plugins__proxy__white_list=["/api/v1/proxy/mock/info","/proxy/mock/get","/proxy/mock/post","/proxy/mock/post_model","/proxy/mock/auth/*"]
     plugins__proxy__auth__general_auth_keys=["i_am_proxy_general_auth_keys"]
     plugins__proxy__auth__auth_urls=["/proxy/mock/auth/*"]
-    plugins__proxy__auth__special_auth_keys={{"/proxy/mock/auth/sget":["i_am_proxy_special_auth_keys"]}}
+    plugins__proxy__auth__special_auth_keys={{"/proxy/mock/auth/sget":["i_am_proxy_special_auth_keys"],"/api/v1/proxy/remote":["i_am_proxy_func_auth_keys"]}}
+    plugins__proxy__proxy_functions={{"http://localhost:9527":["tests.test_plugins.remote_exchange_key_value"]}}
     ; plugins__proxy__force_stream_apis=[]
     sentry__enable=false
     test__silent=true
     auth__general_auth_keys=["i_am_general_auth_keys"]
     auth__auth_urls=["/api/v1/echo"]
+
+
 asyncio_mode = auto

src/framex/__init__.py

@@ -98,7 +98,7 @@ def run(
         from ray import serve
 
         ray.init(
-            num_cpus=num_cpus,
+            num_cpus=num_cpus if num_cpus > 0 else None,
             dashboard_host=dashboard_host,
             dashboard_port=dashboard_port,
             configure_logging=False,

src/framex/adapter/base.py

@@ -29,7 +29,7 @@ def to_deployment(self, cls: type, **kwargs: Any) -> type:  # noqa: ARG002
     async def call_func(self, api: PluginApi, **kwargs: Any) -> Any:
         func = self.get_handle_func(api.deployment_name, api.func_name)
         stream = api.stream
-        if api.call_type == ApiType.PROXY:
+        if api.call_type == ApiType.PROXY and api.api:
             kwargs["proxy_path"] = api.api
             stream = await self._check_is_gen_api(api.api)
         if stream:

src/framex/config.py

@@ -1,4 +1,5 @@
 from typing import Any, Literal, Self
+from uuid import uuid4
 
 from pydantic import BaseModel, Field, model_validator
 from pydantic_settings import (
@@ -9,6 +10,8 @@
     TomlConfigSettingsSource,
 )
 
+from framex.consts import PROXY_FUNC_HTTP_PATH
+
 
 class LogConfig(BaseModel):
     simple_log: bool = True
@@ -41,7 +44,7 @@ class ServerConfig(BaseModel):
     use_ray: bool = False
     enable_proxy: bool = False
     legal_proxy_code: list[int] = [200]
-    num_cpus: int = 8
+    num_cpus: int = -1
     excluded_log_paths: list[str] = []
     ingress_config: dict[str, Any] = {"max_ongoing_requests": 60}
 
@@ -57,7 +60,15 @@ class AuthConfig(BaseModel):
     special_auth_keys: dict[str, list[str]] = Field(default_factory=dict)
 
     @model_validator(mode="after")
-    def validate_special_auth_urls(self) -> Self:
+    def normalize_and_validate(self) -> Self:
+        if PROXY_FUNC_HTTP_PATH not in self.auth_urls:
+            self.auth_urls.append(PROXY_FUNC_HTTP_PATH)
+        if not self.general_auth_keys:  # pragma: no cover
+            from framex.log import logger
+
+            key = str(uuid4())
+            logger.warning(f"No general_auth_keys set, generate a random key: {key}")
+            self.general_auth_keys = [key]
         for special_url in self.special_auth_keys:
             if not self._is_url_protected(special_url):
                 raise ValueError(f"special_auth_keys url '{special_url}' is not covered by any auth_urls rule")

src/framex/consts.py

@@ -4,6 +4,8 @@
 
 BACKEND_NAME = "backend"
 APP_NAME = "default"
+
 PROXY_PLUGIN_NAME = "proxy.ProxyPlugin"
+PROXY_FUNC_HTTP_PATH = f"{API_STR}/proxy/remote"
 
 DEFAULT_ENV = {"RAY_COLOR_PREFIX": "1", "RAY_DEDUP_LOGS": "1", "RAY_SERVE_RUN_SYNC_IN_THREADPOOL": "1"}

src/framex/driver/ingress.py

@@ -44,9 +44,6 @@ def __init__(self, deployments: list[DeploymentHandle], plugin_apis: list["Plugi
                     ApiType.ALL,
                 ]
             ):
-                from framex.config import settings
-
-                auth_keys = settings.auth.get_auth_keys(plugin_api.api)
                 self.register_route(
                     plugin_api.api,
                     plugin_api.methods,
@@ -56,7 +53,6 @@ def __init__(self, deployments: list[DeploymentHandle], plugin_apis: list["Plugi
                     stream=plugin_api.stream,
                     direct_output=False,
                     tags=plugin_api.tags,
-                    auth_keys=auth_keys,
                 )
 
     def register_route(
@@ -71,10 +67,16 @@ def register_route(
         tags: list[str | Enum] | None = None,
         auth_keys: list[str] | None = None,
     ) -> bool:
+        from framex.log import logger
+
         if tags is None:
             tags = ["default"]
+        if auth_keys is None:
+            from framex.config import settings
+
+            auth_keys = settings.auth.get_auth_keys(path)
+            logger.debug(f"API({path}) with tags {tags} requires auth_keys {auth_keys}")
         adapter = get_adapter()
-        from framex.log import logger
 
         try:
             routes: list[str] = [route.path for route in app.routes if isinstance(route, Route | APIRoute)]

src/framex/log.py

@@ -5,8 +5,6 @@
 
 import loguru
 
-from framex.config import settings
-
 if TYPE_CHECKING:
     # avoid sphinx autodoc resolve annotation failed
     # because loguru module do not have `Logger` class actually
@@ -17,6 +15,8 @@
 
 class LoguruHandler(logging.Handler):  # pragma: no cover
     def emit(self, record: logging.LogRecord) -> None:
+        from framex.config import settings
+
         msg = record.getMessage()
         if settings.log.simple_log and (
             (record.name == "ray.serve" and msg.startswith(settings.log.ignored_prefixes)) or record.name == "filelock"

src/framex/plugin/__init__.py

@@ -76,14 +76,17 @@ def init_all_deployments(enable_proxy: bool) -> list[DeploymentHandle]:
 
 
 async def call_plugin_api(
-    api_name: str,
+    api_name: str | PluginApi,
     interval_apis: dict[str, PluginApi] | None = None,
     **kwargs: Any,
 ) -> Any:
-    api = interval_apis.get(api_name) if interval_apis else _manager.get_api(api_name)
+    if isinstance(api_name, PluginApi):
+        api: PluginApi | None = api_name
+    elif isinstance(api_name, str):
+        api = interval_apis.get(api_name) if interval_apis else _manager.get_api(api_name)
     use_proxy = False
     if not api:
-        if api_name.startswith("/") and settings.server.enable_proxy:
+        if isinstance(api_name, str) and api_name.startswith("/") and settings.server.enable_proxy:
             api = PluginApi(
                 api=api_name,
                 deployment_name=PROXY_PLUGIN_NAME,
@@ -116,7 +119,7 @@ async def call_plugin_api(
         return result.model_dump(by_alias=True)
     if use_proxy:
         if not isinstance(result, dict):
-            raise RuntimeError(f"Proxy API {api_name} returned non-dict result: {type(result)}")
+            return result
         if "status" not in result:
             raise RuntimeError(f"Proxy API {api_name} returned invalid response: missing 'status' field")
         res = result.get("data")

src/framex/plugin/load.py

@@ -1,3 +1,5 @@
+from collections.abc import Callable
+
 from framex.config import Settings
 from framex.log import logger
 from framex.plugin.model import Plugin
@@ -36,3 +38,7 @@ def load_from_settings(settings: Settings) -> set[Plugin]:
     builtin_plugin_instances = load_builtin_plugins(*candidate_builtin_plugins) if candidate_builtin_plugins else set()
     plugin_instances = load_plugins(*settings.load_plugins) if settings.load_plugins else set()
     return builtin_plugin_instances | plugin_instances
+
+
+def register_proxy_func(_: Callable) -> None:  # pragma: no cover
+    pass

src/framex/plugin/on.py

@@ -1,3 +1,4 @@
+import functools
 import inspect
 import types
 from collections.abc import Callable
@@ -6,11 +7,11 @@
 from pydantic import BaseModel
 
 from framex.adapter import get_adapter
-from framex.consts import API_STR
+from framex.consts import API_STR, PROXY_PLUGIN_NAME
 from framex.plugin.model import ApiType, PluginApi, PluginDeployment
-from framex.utils import extract_method_params, plugin_to_deployment_name
+from framex.utils import cache_decode, cache_encode, extract_method_params, plugin_to_deployment_name
 
-from . import _current_plugin
+from . import _current_plugin, call_plugin_api
 
 
 def on_register(**kwargs: Any) -> Callable[[type], type]:
@@ -97,6 +98,71 @@ def wrapper(func: Callable) -> Callable:
     return wrapper
 
 
+def on_proxy() -> Callable:
+    def decorator(func: Callable) -> Callable:
+        from framex.config import settings
+
+        if not settings.server.enable_proxy:  # pragma: no cover
+            return func
+
+        is_registered = False
+        full_func_name = f"{func.__module__}.{func.__name__}"
+
+        async def safe_callable(*args: Any, **kwargs: Any) -> Any:
+            try:
+                return await func(*args, **kwargs)
+            except Exception:
+                raw = func
+
+                if isinstance(raw, (classmethod, staticmethod)):
+                    raw = raw.__func__
+
+                while hasattr(raw, "__wrapped__"):
+                    raw = raw.__wrapped__
+
+                if inspect.iscoroutinefunction(raw):
+                    return await raw(*args, **kwargs)
+                return raw(*args, **kwargs)
+
+        @functools.wraps(func)
+        async def wrapper(*args: Any, **kwargs: Any) -> Any:
+            nonlocal is_registered
+
+            if args:  # pragma: no cover
+                raise TypeError(f"The proxy function '{func.__name__}' only supports keyword arguments.")
+
+            if not is_registered:
+                api_reg = PluginApi(
+                    deployment_name=PROXY_PLUGIN_NAME,
+                    call_type=ApiType.PROXY,
+                    func_name="register_proxy_function",
+                )
+                await call_plugin_api(
+                    api_reg,
+                    None,
+                    func_name=full_func_name,
+                    func_callable=safe_callable,
+                )
+                is_registered = True
+
+            api_call = PluginApi(
+                deployment_name=PROXY_PLUGIN_NAME,
+                call_type=ApiType.PROXY,
+                func_name="call_proxy_function",
+            )
+            res = await call_plugin_api(
+                api_call,
+                None,
+                func_name=cache_encode(full_func_name),
+                data=cache_encode(data=kwargs),
+            )
+            return cache_decode(res)
+
+        return wrapper
+
+    return decorator
+
+
 def remote() -> Callable:
     def wrapper(func: Callable) -> Any:
         adapter = get_adapter()