Skip to content

patcher

agent_cover.instrumentation.agents.patcher

Instrumentation module for agent runners.

This module provides the main instrumentor class responsible for applying wrappers to agent execution methods.

Context Tracking

This module wraps methods like AgentExecutor.invoke to set a "Context Token". This allows AgentCover to know when an LLM call is happening inside an agent versus outside (e.g., during test setup), ensuring accurate metrics.

Attributes

Classes

AgentInstrumentor

Bases: BaseInstrumentor

Instruments agent classes with context management strategies.

This class handles the lifecycle of patching agent runners. It identifies target modules, loads them if necessary, resolves the specific classes, and wraps execution methods with the appropriate context managers.

Attributes:

Name Type Description
targets_provider

A callable returning a list of targets to instrument.
strategies

A dictionary mapping strategy keys to WrapperStrategy instances.
registry

The agent registry instance.
context_manager

The context manager for handling agent contexts.
patch_manager

The manager responsible for applying and reverting patches.
module_iterator

A callable returning the current snapshot of sys.modules.
importer_func

A callable used to import modules by name.
is_instrumented

A boolean indicating if instrumentation has been applied.

Methods:

Name Description
instrument

Performs the instrumentation process on all defined targets.
Source code in src/agent_cover/instrumentation/agents/patcher.py 
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
class AgentInstrumentor(BaseInstrumentor):
    """Instruments agent classes with context management strategies.

    This class handles the lifecycle of patching agent runners. It identifies
    target modules, loads them if necessary, resolves the specific classes,
    and wraps execution methods with the appropriate context managers.

    Attributes:
        targets_provider: A callable returning a list of targets to instrument.
        strategies: A dictionary mapping strategy keys to WrapperStrategy instances.
        registry: The agent registry instance.
        context_manager: The context manager for handling agent contexts.
        patch_manager: The manager responsible for applying and reverting patches.
        module_iterator: A callable returning the current snapshot of sys.modules.
        importer_func: A callable used to import modules by name.
        is_instrumented: A boolean indicating if instrumentation has been applied.

    Methods:
        instrument: Performs the instrumentation process on all defined targets.
    """

    def __init__(
        self,
        registry: Optional[AgentRegistry] = None,
        context_manager: Optional[AgentContextManager] = None,
        patch_manager: Optional[PatchManager] = None,
        module_iterator: Optional[Callable[[], Dict[str, Any]]] = None,
        importer_func: Optional[Callable[[str], Any]] = None,
        targets_provider: Optional[Callable[[], TargetList]] = None,
        strategies: Optional[Dict[str, WrapperStrategy]] = None,
    ):
        """Initializes the AgentInstrumentor.

        Args:
            registry: Optional registry for tracking agents.
            context_manager: Optional manager for agent context.
            patch_manager: Optional manager for handling patches.
            module_iterator: Optional callable to get current modules.
            importer_func: Optional callable to import modules.
            targets_provider: Optional callable to provide instrumentation targets.
            strategies: Optional dictionary of instrumentation strategies.
        """
        super().__init__(
            registry, context_manager, patch_manager, module_iterator, importer_func
        )
        self.targets_provider = targets_provider or _default_targets_provider
        self.strategies = strategies or DEFAULT_STRATEGIES

    def instrument(self):
        """Applies instrumentation to all discoverable agent targets.

        This method iterates through the targets provided by the targets_provider.
        It attempts to load the module if it is not present, resolves the target
        class using the base class resolver, and applies the configured wrapper
        strategies to the specified methods.
        """
        if self.is_instrumented:
            return

        raw_targets = self.targets_provider()
        targets = self._normalize_targets(raw_targets)

        modules_snapshot = self.module_iterator()

        for target in targets:
            if not self._should_instrument(target):
                continue

            mod_name = target.module
            cls_name = target.class_name
            method_map = target.methods

            try:
                # Attempt lazy import
                if mod_name not in modules_snapshot:
                    try:
                        self.importer(mod_name)
                        modules_snapshot = self.module_iterator()
                    except ModuleNotFoundError:
                        Log.log_skip_missing_module(logger, mod_name)
                        continue
                    except Exception as e:
                        logger.warning(e, exc_info=True)
                        continue

                if mod_name in modules_snapshot:
                    mod = modules_snapshot[mod_name]

                    # [CHANGE] Use the robust resolver from the base class
                    cls = self._resolve_target_class(mod, cls_name)

                    if cls and "_ac_agent_patched" not in cls.__dict__:
                        self._patch_agent_runner(cls, method_map)
                        cls._ac_agent_patched = True

            except Exception as e:
                logger.warning(e, exc_info=True)

        self.is_instrumented = True

    def _patch_agent_runner(self, cls: Any, method_map: Dict[str, str]):
        """Applies wrappers to the specified methods of an agent class.

        Args:
            cls: The class object to patch.
            method_map: A dictionary mapping method names to strategy keys.
        """
        for method_name, strategy_key in method_map.items():
            if not hasattr(cls, method_name):
                continue

            strategy = self.strategies.get(strategy_key)
            if not strategy:
                continue

            original = getattr(cls, method_name)
            wrapper = strategy.wrap(original, self.context_manager)

            self._safe_patch(cls, method_name, wrapper)
Functions
__init__(registry=None, context_manager=None, patch_manager=None, module_iterator=None, importer_func=None, targets_provider=None, strategies=None)

Initializes the AgentInstrumentor.

Parameters:

Name Type Description Default
registry Optional[AgentRegistry]

Optional registry for tracking agents.

 None
context_manager Optional[AgentContextManager]

Optional manager for agent context.

 None
patch_manager Optional[PatchManager]

Optional manager for handling patches.

 None
module_iterator Optional[Callable[[], Dict[str, Any]]]

Optional callable to get current modules.

 None
importer_func Optional[Callable[[str], Any]]

Optional callable to import modules.

 None
targets_provider Optional[Callable[[], TargetList]]

Optional callable to provide instrumentation targets.

 None
strategies Optional[Dict[str, WrapperStrategy]]

Optional dictionary of instrumentation strategies.

 None
Source code in src/agent_cover/instrumentation/agents/patcher.py 
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
def __init__(
    self,
    registry: Optional[AgentRegistry] = None,
    context_manager: Optional[AgentContextManager] = None,
    patch_manager: Optional[PatchManager] = None,
    module_iterator: Optional[Callable[[], Dict[str, Any]]] = None,
    importer_func: Optional[Callable[[str], Any]] = None,
    targets_provider: Optional[Callable[[], TargetList]] = None,
    strategies: Optional[Dict[str, WrapperStrategy]] = None,
):
    """Initializes the AgentInstrumentor.

    Args:
        registry: Optional registry for tracking agents.
        context_manager: Optional manager for agent context.
        patch_manager: Optional manager for handling patches.
        module_iterator: Optional callable to get current modules.
        importer_func: Optional callable to import modules.
        targets_provider: Optional callable to provide instrumentation targets.
        strategies: Optional dictionary of instrumentation strategies.
    """
    super().__init__(
        registry, context_manager, patch_manager, module_iterator, importer_func
    )
    self.targets_provider = targets_provider or _default_targets_provider
    self.strategies = strategies or DEFAULT_STRATEGIES
instrument()

Applies instrumentation to all discoverable agent targets.

This method iterates through the targets provided by the targets_provider. It attempts to load the module if it is not present, resolves the target class using the base class resolver, and applies the configured wrapper strategies to the specified methods.

Source code in src/agent_cover/instrumentation/agents/patcher.py 
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
def instrument(self):
    """Applies instrumentation to all discoverable agent targets.

    This method iterates through the targets provided by the targets_provider.
    It attempts to load the module if it is not present, resolves the target
    class using the base class resolver, and applies the configured wrapper
    strategies to the specified methods.
    """
    if self.is_instrumented:
        return

    raw_targets = self.targets_provider()
    targets = self._normalize_targets(raw_targets)

    modules_snapshot = self.module_iterator()

    for target in targets:
        if not self._should_instrument(target):
            continue

        mod_name = target.module
        cls_name = target.class_name
        method_map = target.methods

        try:
            # Attempt lazy import
            if mod_name not in modules_snapshot:
                try:
                    self.importer(mod_name)
                    modules_snapshot = self.module_iterator()
                except ModuleNotFoundError:
                    Log.log_skip_missing_module(logger, mod_name)
                    continue
                except Exception as e:
                    logger.warning(e, exc_info=True)
                    continue

            if mod_name in modules_snapshot:
                mod = modules_snapshot[mod_name]

                # [CHANGE] Use the robust resolver from the base class
                cls = self._resolve_target_class(mod, cls_name)

                if cls and "_ac_agent_patched" not in cls.__dict__:
                    self._patch_agent_runner(cls, method_map)
                    cls._ac_agent_patched = True

        except Exception as e:
            logger.warning(e, exc_info=True)

    self.is_instrumented = True

Functions

instrument_agents(registry=None)

Helper function to instantiate and run the AgentInstrumentor.

This function exists primarily for backward compatibility.

Parameters:

Name Type Description Default
registry Optional[AgentRegistry]

An optional AgentRegistry instance.

 None

Returns:

Type Description
AgentInstrumentor

The initialized and executed AgentInstrumentor instance.
Source code in src/agent_cover/instrumentation/agents/patcher.py 
173
174
175
176
177
178
179
180
181
182
183
184
185
186
def instrument_agents(registry: Optional[AgentRegistry] = None) -> AgentInstrumentor:
    """Helper function to instantiate and run the AgentInstrumentor.

    This function exists primarily for backward compatibility.

    Args:
        registry: An optional AgentRegistry instance.

    Returns:
        The initialized and executed AgentInstrumentor instance.
    """
    instrumentor = AgentInstrumentor(registry=registry)
    instrumentor.instrument()
    return instrumentor