fix(utils): Preserve decorated type for @experimental and @working_in_progress

src/google/adk/utils/feature_decorator.py

@@ -20,12 +20,35 @@
 from typing import Any
 from typing import cast
 from typing import Optional
+from typing import overload
+from typing import Protocol
 from typing import TypeVar
 import warnings
 
 T = TypeVar("T")
 
 
+class _FeatureDecorator(Protocol):
+  """A feature decorator usable with or without a message argument.
+
+  Preserves the decorated object's type so that subclasses and type
+  checkers continue to see the real class/function rather than ``Any``.
+  """
+
+  # @decorator (bare, on a class or function)
+  @overload
+  def __call__(self, message_or_obj: T) -> T:
+    ...
+
+  # @decorator() or @decorator("message")
+  @overload
+  def __call__(self, message_or_obj: Optional[str] = ...) -> Callable[[T], T]:
+    ...
+
+  def __call__(self, message_or_obj: Any = None) -> Any:
+    ...
+
+
 def _is_truthy_env(var_name: str) -> bool:
   value = os.environ.get(var_name)
   if value is None:
@@ -39,7 +62,7 @@ def _make_feature_decorator(
     default_message: str,
     block_usage: bool = False,
     bypass_env_var: Optional[str] = None,
-) -> Callable[..., Any]:
+) -> _FeatureDecorator:
   def decorator_factory(message_or_obj: Any = None) -> Any:
     # Case 1: Used as @decorator without parentheses
     # message_or_obj is the decorated class/function
@@ -57,7 +80,7 @@ def decorator_factory(message_or_obj: Any = None) -> Any:
     )
     return _create_decorator(message, label, block_usage, bypass_env_var)
 
-  return decorator_factory
+  return cast(_FeatureDecorator, decorator_factory)
 
 
 def _create_decorator(