From 37b217751b555a840845ead1e288c9bac4181da8 Mon Sep 17 00:00:00 2001 From: Chaitanya Laxman Date: Mon, 7 Sep 2026 18:35:50 +0400 Subject: [PATCH 1/3] fix(plugins): wrap class descriptors from __dict__ getmembers unwraps staticmethod and skips classmethod, so instrumentation changed call semantics and pinned base methods onto subclasses. Walk the class dict and put the same descriptor back. Fixes #6980 --- src/google/adk/plugins/auto_tracing_plugin.py | 56 ++++++++++-- .../plugins/test_auto_tracing_plugin.py | 90 +++++++++++++++++++ 2 files changed, 138 insertions(+), 8 deletions(-) diff --git a/src/google/adk/plugins/auto_tracing_plugin.py b/src/google/adk/plugins/auto_tracing_plugin.py index df993524b0f..a99bca29edb 100644 --- a/src/google/adk/plugins/auto_tracing_plugin.py +++ b/src/google/adk/plugins/auto_tracing_plugin.py @@ -152,14 +152,54 @@ def _wrap_module(self, module: ModuleType) -> None: if inspect.isfunction(attr): self._rebind(module, attr_name, attr) elif inspect.isclass(attr): - for member_name, member in inspect.getmembers(attr): - if member_name.startswith("__"): - continue - if not inspect.isfunction(member): - continue - if getattr(member, "__module__", "") != module_name: - continue - self._rebind(attr, member_name, member) + self._wrap_class(attr, module_name) + + def _wrap_class(self, cls: type[Any], module_name: str) -> None: + """Wraps callables defined on ``cls``, not inherited unwrapped members.""" + for member_name, member in cls.__dict__.items(): + if member_name.startswith("__"): + continue + if isinstance(member, staticmethod): + fn = member.__func__ + if getattr(fn, "__module__", "") != module_name: + continue + self._rebind_descriptor(cls, member_name, fn, staticmethod) + elif isinstance(member, classmethod): + fn = member.__func__ + if getattr(fn, "__module__", "") != module_name: + continue + self._rebind_descriptor(cls, member_name, fn, classmethod) + elif inspect.isfunction(member): + if getattr(member, "__module__", "") != module_name: + continue + self._rebind(cls, member_name, member) + + def _rebind_descriptor( + self, + owner: type[Any], + name: str, + fn: Callable[..., Any], + descriptor: type[staticmethod] | type[classmethod], + ) -> None: + if getattr(fn, auto_tracing_helpers.WRAPPED_ATTR, False): + return + try: + setattr( + owner, + name, + descriptor( + auto_tracing_helpers.build_tracing_wrapper( + fn, self._tracer, self._caps + ) + ), + ) + except (AttributeError, TypeError) as exc: + logger.info( + "AutoTracingPlugin: cannot rebind %s.%s: %s", + getattr(owner, "__qualname__", owner), + name, + exc, + ) def _rebind( self, owner: ModuleType | type[Any], name: str, fn: Callable[..., Any] diff --git a/tests/unittests/plugins/test_auto_tracing_plugin.py b/tests/unittests/plugins/test_auto_tracing_plugin.py index 3e5dd27e1be..76ed9d77917 100644 --- a/tests/unittests/plugins/test_auto_tracing_plugin.py +++ b/tests/unittests/plugins/test_auto_tracing_plugin.py @@ -752,3 +752,93 @@ def producer(): assert f"first {cap}:" in rendered, rendered finally: sys.modules.pop(name, None) + + +_DESCRIPTOR_MODULE_NAME = ( + "google.adk.tests.unittests.plugins.descriptor_test_fixture" +) + + +def _build_descriptor_module() -> types.ModuleType: + module = types.ModuleType(_DESCRIPTOR_MODULE_NAME) + module.__name__ = _DESCRIPTOR_MODULE_NAME + + def slugify(text): + return text.strip().lower().replace(" ", "-") + + def build(cls, name): + return cls.slugify(name) + + def instance_method(self, x): + return x + 1 + + def shared(self, x): + return x * 2 + + for fn in (slugify, build, instance_method, shared): + fn.__module__ = _DESCRIPTOR_MODULE_NAME + + tools = type( + "Tools", + (), + { + "slugify": staticmethod(slugify), + "build": classmethod(build), + "instance_method": instance_method, + }, + ) + zbase = type("ZBase", (), {"shared": shared}) + achild = type("AChild", (zbase,), {}) + for cls in (tools, zbase, achild): + cls.__module__ = _DESCRIPTOR_MODULE_NAME + module.AChild = achild + module.Tools = tools + module.ZBase = zbase + return module + + +def test_staticmethod_stays_callable_on_instance(fixture): + module = _build_descriptor_module() + sys.modules[_DESCRIPTOR_MODULE_NAME] = module + try: + plugin = auto_tracing_plugin.AutoTracingPlugin( + tracer=fixture.tracer, + extra_scope_prefixes=(_DESCRIPTOR_MODULE_NAME,), + ) + asyncio.run(plugin.before_run_callback(invocation_context=None)) + assert isinstance(module.Tools.__dict__["slugify"], staticmethod) + assert module.Tools().slugify("Hello World") == "hello-world" + assert any("slugify" in n for n in _span_names(fixture.exporter)) + finally: + sys.modules.pop(_DESCRIPTOR_MODULE_NAME, None) + + +def test_classmethod_is_traced(fixture): + module = _build_descriptor_module() + sys.modules[_DESCRIPTOR_MODULE_NAME] = module + try: + plugin = auto_tracing_plugin.AutoTracingPlugin( + tracer=fixture.tracer, + extra_scope_prefixes=(_DESCRIPTOR_MODULE_NAME,), + ) + asyncio.run(plugin.before_run_callback(invocation_context=None)) + assert isinstance(module.Tools.__dict__["build"], classmethod) + assert module.Tools.build("Hello World") == "hello-world" + assert any("build" in n for n in _span_names(fixture.exporter)) + finally: + sys.modules.pop(_DESCRIPTOR_MODULE_NAME, None) + + +def test_inherited_method_is_not_pinned_on_subclass(fixture): + module = _build_descriptor_module() + sys.modules[_DESCRIPTOR_MODULE_NAME] = module + try: + plugin = auto_tracing_plugin.AutoTracingPlugin( + tracer=fixture.tracer, + extra_scope_prefixes=(_DESCRIPTOR_MODULE_NAME,), + ) + asyncio.run(plugin.before_run_callback(invocation_context=None)) + assert "shared" not in module.AChild.__dict__ + assert module.AChild().shared(3) == 6 + finally: + sys.modules.pop(_DESCRIPTOR_MODULE_NAME, None) From 81553589830a90028cf48f8c04557217dec06d77 Mon Sep 17 00:00:00 2001 From: Chaitanya Laxman Date: Mon, 7 Sep 2026 18:43:22 +0400 Subject: [PATCH 2/3] fix(plugins): fold descriptor wrap into _rebind Avoid a second rebind helper and type[staticmethod] mypy errors. Descriptor factories go through an optional wrap callable. Fixes #6980 --- src/google/adk/plugins/auto_tracing_plugin.py | 76 +++++++------------ .../plugins/test_auto_tracing_plugin.py | 1 + 2 files changed, 29 insertions(+), 48 deletions(-) diff --git a/src/google/adk/plugins/auto_tracing_plugin.py b/src/google/adk/plugins/auto_tracing_plugin.py index a99bca29edb..efabe2620c8 100644 --- a/src/google/adk/plugins/auto_tracing_plugin.py +++ b/src/google/adk/plugins/auto_tracing_plugin.py @@ -155,65 +155,45 @@ def _wrap_module(self, module: ModuleType) -> None: self._wrap_class(attr, module_name) def _wrap_class(self, cls: type[Any], module_name: str) -> None: - """Wraps callables defined on ``cls``, not inherited unwrapped members.""" - for member_name, member in cls.__dict__.items(): + """Wraps functions, staticmethods and classmethods in cls.__dict__. + + Inherited members are left to the defining class so they are not + pinned onto subclasses. + """ + for member_name, member in list(cls.__dict__.items()): if member_name.startswith("__"): continue - if isinstance(member, staticmethod): - fn = member.__func__ - if getattr(fn, "__module__", "") != module_name: - continue - self._rebind_descriptor(cls, member_name, fn, staticmethod) - elif isinstance(member, classmethod): - fn = member.__func__ - if getattr(fn, "__module__", "") != module_name: - continue - self._rebind_descriptor(cls, member_name, fn, classmethod) - elif inspect.isfunction(member): - if getattr(member, "__module__", "") != module_name: - continue - self._rebind(cls, member_name, member) - - def _rebind_descriptor( - self, - owner: type[Any], - name: str, - fn: Callable[..., Any], - descriptor: type[staticmethod] | type[classmethod], - ) -> None: - if getattr(fn, auto_tracing_helpers.WRAPPED_ATTR, False): - return - try: - setattr( - owner, - name, - descriptor( - auto_tracing_helpers.build_tracing_wrapper( - fn, self._tracer, self._caps - ) - ), + fn = ( + member.__func__ + if isinstance(member, (staticmethod, classmethod)) + else member ) - except (AttributeError, TypeError) as exc: - logger.info( - "AutoTracingPlugin: cannot rebind %s.%s: %s", - getattr(owner, "__qualname__", owner), - name, - exc, + if ( + not inspect.isfunction(fn) + or getattr(fn, "__module__", "") != module_name + ): + continue + self._rebind( + cls, + member_name, + fn, + wrap=type(member) if fn is not member else None, ) def _rebind( - self, owner: ModuleType | type[Any], name: str, fn: Callable[..., Any] + self, + owner: ModuleType | type[Any], + name: str, + fn: Callable[..., Any], + wrap: Callable[[Callable[..., Any]], object] | None = None, ) -> None: if getattr(fn, auto_tracing_helpers.WRAPPED_ATTR, False): return try: - setattr( - owner, - name, - auto_tracing_helpers.build_tracing_wrapper( - fn, self._tracer, self._caps - ), + wrapper = auto_tracing_helpers.build_tracing_wrapper( + fn, self._tracer, self._caps ) + setattr(owner, name, wrap(wrapper) if wrap else wrapper) except (AttributeError, TypeError) as exc: logger.info( "AutoTracingPlugin: cannot rebind %s.%s: %s", diff --git a/tests/unittests/plugins/test_auto_tracing_plugin.py b/tests/unittests/plugins/test_auto_tracing_plugin.py index 76ed9d77917..838cfbfd717 100644 --- a/tests/unittests/plugins/test_auto_tracing_plugin.py +++ b/tests/unittests/plugins/test_auto_tracing_plugin.py @@ -840,5 +840,6 @@ def test_inherited_method_is_not_pinned_on_subclass(fixture): asyncio.run(plugin.before_run_callback(invocation_context=None)) assert "shared" not in module.AChild.__dict__ assert module.AChild().shared(3) == 6 + assert any("shared" in n for n in _span_names(fixture.exporter)) finally: sys.modules.pop(_DESCRIPTOR_MODULE_NAME, None) From 313b94411ad74aafd25462bce7e2352dfee84255 Mon Sep 17 00:00:00 2001 From: Chaitanya Laxman Date: Wed, 9 Sep 2026 03:33:06 +0400 Subject: [PATCH 3/3] test(plugins): lock async staticmethod and classmethod wrap Sync descriptor tests stay green if _rebind drops the wrap factory. Async staticmethod unwraps and async classmethod emits no span. Fixes #6980 --- .../plugins/test_auto_tracing_plugin.py | 49 ++++++++++++++++++- 1 file changed, 48 insertions(+), 1 deletion(-) diff --git a/tests/unittests/plugins/test_auto_tracing_plugin.py b/tests/unittests/plugins/test_auto_tracing_plugin.py index 838cfbfd717..bf009a080d7 100644 --- a/tests/unittests/plugins/test_auto_tracing_plugin.py +++ b/tests/unittests/plugins/test_auto_tracing_plugin.py @@ -775,7 +775,20 @@ def instance_method(self, x): def shared(self, x): return x * 2 - for fn in (slugify, build, instance_method, shared): + async def async_slugify(text): + return text.strip().lower().replace(" ", "-") + + async def async_build(cls, name): + return await cls.async_slugify(name) + + for fn in ( + slugify, + build, + instance_method, + shared, + async_slugify, + async_build, + ): fn.__module__ = _DESCRIPTOR_MODULE_NAME tools = type( @@ -785,6 +798,8 @@ def shared(self, x): "slugify": staticmethod(slugify), "build": classmethod(build), "instance_method": instance_method, + "async_slugify": staticmethod(async_slugify), + "async_build": classmethod(async_build), }, ) zbase = type("ZBase", (), {"shared": shared}) @@ -843,3 +858,35 @@ def test_inherited_method_is_not_pinned_on_subclass(fixture): assert any("shared" in n for n in _span_names(fixture.exporter)) finally: sys.modules.pop(_DESCRIPTOR_MODULE_NAME, None) + + +async def test_async_staticmethod_stays_callable_on_instance(fixture): + module = _build_descriptor_module() + sys.modules[_DESCRIPTOR_MODULE_NAME] = module + try: + plugin = auto_tracing_plugin.AutoTracingPlugin( + tracer=fixture.tracer, + extra_scope_prefixes=(_DESCRIPTOR_MODULE_NAME,), + ) + await plugin.before_run_callback(invocation_context=None) + assert isinstance(module.Tools.__dict__["async_slugify"], staticmethod) + assert await module.Tools().async_slugify("Hello World") == "hello-world" + assert any("async_slugify" in n for n in _span_names(fixture.exporter)) + finally: + sys.modules.pop(_DESCRIPTOR_MODULE_NAME, None) + + +async def test_async_classmethod_is_traced(fixture): + module = _build_descriptor_module() + sys.modules[_DESCRIPTOR_MODULE_NAME] = module + try: + plugin = auto_tracing_plugin.AutoTracingPlugin( + tracer=fixture.tracer, + extra_scope_prefixes=(_DESCRIPTOR_MODULE_NAME,), + ) + await plugin.before_run_callback(invocation_context=None) + assert isinstance(module.Tools.__dict__["async_build"], classmethod) + assert await module.Tools.async_build("Hello World") == "hello-world" + assert any("async_build" in n for n in _span_names(fixture.exporter)) + finally: + sys.modules.pop(_DESCRIPTOR_MODULE_NAME, None)