diff --git a/src/google/adk/plugins/auto_tracing_plugin.py b/src/google/adk/plugins/auto_tracing_plugin.py index df993524b0f..efabe2620c8 100644 --- a/src/google/adk/plugins/auto_tracing_plugin.py +++ b/src/google/adk/plugins/auto_tracing_plugin.py @@ -152,28 +152,48 @@ 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 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 + fn = ( + member.__func__ + if isinstance(member, (staticmethod, classmethod)) + else member + ) + 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 3e5dd27e1be..bf009a080d7 100644 --- a/tests/unittests/plugins/test_auto_tracing_plugin.py +++ b/tests/unittests/plugins/test_auto_tracing_plugin.py @@ -752,3 +752,141 @@ 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 + + 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( + "Tools", + (), + { + "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}) + 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 + 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)