Skip to content

strategies

agent_cover.instrumentation.strategies

Defines wrapper strategies for agent execution context management.

This module provides abstract and concrete implementations of wrapper strategies used to intercept method calls. These strategies are primarily responsible for managing the execution context (activating and resetting tokens) during synchronous, asynchronous, and generator executions within the agent coverage framework.

Classes

AsyncContextWrapper

Bases: BaseExecutionStrategy

Wraps asynchronous methods (coroutines) to maintain Agent context.

This is critical for frameworks like LangChain where execution jumps between threads and event loops. It ensures that the contextvars token is set correctly before awaiting the original coroutine and reset immediately after.

Parameters:

Name Type Description Default
original

The original async def function.

 required
ctx_manager

The manager handling the context token.

 required

Returns:

Name Type Description
Callable

An async wrapper that maintains scope.
Source code in src/agent_cover/instrumentation/strategies.py 
 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
class AsyncContextWrapper(BaseExecutionStrategy):
    """Wraps asynchronous methods (coroutines) to maintain Agent context.

    This is critical for frameworks like LangChain where execution jumps between
    threads and event loops. It ensures that the `contextvars` token is set
    correctly before `await`ing the original coroutine and reset immediately after.

    Args:
        original: The original `async def` function.
        ctx_manager: The manager handling the context token.

    Returns:
        Callable: An async wrapper that maintains scope.
    """

    def wrap(self, original: Callable, ctx_manager: AgentContextManager) -> Callable:
        """Wraps an asynchronous method to activate and deactivate the agent context.

        Args:
            original: The original asynchronous method.
            ctx_manager: The AgentContextManager used to manage the execution context.

        Returns:
            Callable: An asynchronous wrapper function that manages the agent context.
        """

        async def wrapper(instance: Any, *args: Any, **kwargs: Any) -> Any:
            """The asynchronous wrapper function handling the context lifecycle."""
            token = ctx_manager.set_active(True)  # Activate the agent context
            try:
                return await original(
                    instance, *args, **kwargs
                )  # Execute the original method
            finally:
                ctx_manager.reset(token)  # Reset the agent context

        return wrapper
Functions
wrap(original, ctx_manager)

Wraps an asynchronous method to activate and deactivate the agent context.

Parameters:

Name Type Description Default
original Callable

The original asynchronous method.

 required
ctx_manager AgentContextManager

The AgentContextManager used to manage the execution context.

 required

Returns:

Name Type Description
Callable Callable

An asynchronous wrapper function that manages the agent context.
Source code in src/agent_cover/instrumentation/strategies.py 
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
def wrap(self, original: Callable, ctx_manager: AgentContextManager) -> Callable:
    """Wraps an asynchronous method to activate and deactivate the agent context.

    Args:
        original: The original asynchronous method.
        ctx_manager: The AgentContextManager used to manage the execution context.

    Returns:
        Callable: An asynchronous wrapper function that manages the agent context.
    """

    async def wrapper(instance: Any, *args: Any, **kwargs: Any) -> Any:
        """The asynchronous wrapper function handling the context lifecycle."""
        token = ctx_manager.set_active(True)  # Activate the agent context
        try:
            return await original(
                instance, *args, **kwargs
            )  # Execute the original method
        finally:
            ctx_manager.reset(token)  # Reset the agent context

    return wrapper

AsyncGenContextWrapper

Bases: BaseExecutionStrategy

Strategy for wrapping asynchronous generator methods to manage the agent context.

Similar to the sync generator, we must ensure context is active only during the await / processing phases inside the generator, and closed during yield.

Source code in src/agent_cover/instrumentation/strategies.py 
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
class AsyncGenContextWrapper(BaseExecutionStrategy):
    """Strategy for wrapping asynchronous generator methods to manage the agent context.

    Similar to the sync generator, we must ensure context is active only during
    the `await` / processing phases inside the generator, and closed during yield.
    """

    def wrap(self, original: Callable, ctx_manager: AgentContextManager) -> Callable:
        """Wraps an asynchronous generator method.

        Manually iterates using `__anext__` to control context boundaries.
        """

        async def wrapper(
            instance: Any, *args: Any, **kwargs: Any
        ) -> AsyncIterator[Any]:
            iterator = original(instance, *args, **kwargs)
            while True:
                token = ctx_manager.set_active(True)
                try:
                    # Use manual iteration protocol for async generators
                    value = await iterator.__anext__()
                except StopAsyncIteration:
                    ctx_manager.reset(token)
                    return
                except Exception as e:
                    logger.warning(e, exc_info=True)
                    ctx_manager.reset(token)
                    raise

                # Reset BEFORE yielding to caller
                ctx_manager.reset(token)
                yield value

        return wrapper
Functions
wrap(original, ctx_manager)

Wraps an asynchronous generator method.

Manually iterates using __anext__ to control context boundaries.

Source code in src/agent_cover/instrumentation/strategies.py 
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
def wrap(self, original: Callable, ctx_manager: AgentContextManager) -> Callable:
    """Wraps an asynchronous generator method.

    Manually iterates using `__anext__` to control context boundaries.
    """

    async def wrapper(
        instance: Any, *args: Any, **kwargs: Any
    ) -> AsyncIterator[Any]:
        iterator = original(instance, *args, **kwargs)
        while True:
            token = ctx_manager.set_active(True)
            try:
                # Use manual iteration protocol for async generators
                value = await iterator.__anext__()
            except StopAsyncIteration:
                ctx_manager.reset(token)
                return
            except Exception as e:
                logger.warning(e, exc_info=True)
                ctx_manager.reset(token)
                raise

            # Reset BEFORE yielding to caller
            ctx_manager.reset(token)
            yield value

    return wrapper

BaseExecutionStrategy

Bases: WrapperStrategy

Base class for strategies that require the ContextManager.

Methods:

Name Description
wrap

Inherited from WrapperStrategy.
Source code in src/agent_cover/instrumentation/strategies.py 
47
48
49
50
51
52
53
54
55
56
57
class BaseExecutionStrategy(WrapperStrategy):
    """Base class for strategies that require the ContextManager.

    Attributes:
        None

    Methods:
        wrap: Inherited from WrapperStrategy.
    """

    pass

SyncContextWrapper

Bases: BaseExecutionStrategy

Strategy for wrapping synchronous methods to manage the agent context.

Methods:

Name Description
wrap

Wraps a synchronous callable to handle context activation/deactivation.
Source code in src/agent_cover/instrumentation/strategies.py 
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
class SyncContextWrapper(BaseExecutionStrategy):
    """Strategy for wrapping synchronous methods to manage the agent context.

    Attributes:
        None

    Methods:
        wrap: Wraps a synchronous callable to handle context activation/deactivation.
    """

    def wrap(self, original: Callable, ctx_manager: AgentContextManager) -> Callable:
        """Wraps a synchronous method to activate and deactivate the agent context.

        Args:
            original: The original synchronous method.
            ctx_manager: The AgentContextManager used to manage the execution context.

        Returns:
            Callable: A wrapper function that manages the agent context around the original call.
        """

        def wrapper(instance: Any, *args: Any, **kwargs: Any) -> Any:
            """The wrapper function handling the context lifecycle."""
            token = ctx_manager.set_active(True)  # Activate the agent context
            try:
                return original(
                    instance, *args, **kwargs
                )  # Execute the original method
            finally:
                ctx_manager.reset(token)  # Reset the agent context

        return wrapper
Functions
wrap(original, ctx_manager)

Wraps a synchronous method to activate and deactivate the agent context.

Parameters:

Name Type Description Default
original Callable

The original synchronous method.

 required
ctx_manager AgentContextManager

The AgentContextManager used to manage the execution context.

 required

Returns:

Name Type Description
Callable Callable

A wrapper function that manages the agent context around the original call.
Source code in src/agent_cover/instrumentation/strategies.py 
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
def wrap(self, original: Callable, ctx_manager: AgentContextManager) -> Callable:
    """Wraps a synchronous method to activate and deactivate the agent context.

    Args:
        original: The original synchronous method.
        ctx_manager: The AgentContextManager used to manage the execution context.

    Returns:
        Callable: A wrapper function that manages the agent context around the original call.
    """

    def wrapper(instance: Any, *args: Any, **kwargs: Any) -> Any:
        """The wrapper function handling the context lifecycle."""
        token = ctx_manager.set_active(True)  # Activate the agent context
        try:
            return original(
                instance, *args, **kwargs
            )  # Execute the original method
        finally:
            ctx_manager.reset(token)  # Reset the agent context

    return wrapper

SyncGenContextWrapper

Bases: BaseExecutionStrategy

Wraps synchronous generators (functions using yield).

Generators are tricky because execution pauses and resumes. This wrapper ensures the Agent context is active every time the generator executes code, but inactive when it yields control back to the caller.

Source code in src/agent_cover/instrumentation/strategies.py 
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
class SyncGenContextWrapper(BaseExecutionStrategy):
    """Wraps synchronous generators (functions using `yield`).

    Generators are tricky because execution pauses and resumes. This wrapper
    ensures the Agent context is active *every time* the generator executes code,
    but inactive when it yields control back to the caller.
    """

    def wrap(self, original: Callable, ctx_manager: AgentContextManager) -> Callable:
        """Wraps a synchronous generator method.

        We cannot use `yield from` inside a `try/finally` because that keeps
        the context active while the generator is suspended. We must iterate manually.
        """

        def wrapper(instance: Any, *args: Any, **kwargs: Any) -> Iterator[Any]:
            iterator = original(instance, *args, **kwargs)
            while True:
                token = ctx_manager.set_active(True)
                try:
                    value = next(iterator)
                except StopIteration:
                    ctx_manager.reset(token)
                    return
                except Exception as e:
                    logger.warning(e, exc_info=True)
                    ctx_manager.reset(token)
                    raise

                # Reset BEFORE yielding to caller
                ctx_manager.reset(token)
                yield value

        return wrapper
Functions
wrap(original, ctx_manager)

Wraps a synchronous generator method.

We cannot use yield from inside a try/finally because that keeps the context active while the generator is suspended. We must iterate manually.

Source code in src/agent_cover/instrumentation/strategies.py 
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
def wrap(self, original: Callable, ctx_manager: AgentContextManager) -> Callable:
    """Wraps a synchronous generator method.

    We cannot use `yield from` inside a `try/finally` because that keeps
    the context active while the generator is suspended. We must iterate manually.
    """

    def wrapper(instance: Any, *args: Any, **kwargs: Any) -> Iterator[Any]:
        iterator = original(instance, *args, **kwargs)
        while True:
            token = ctx_manager.set_active(True)
            try:
                value = next(iterator)
            except StopIteration:
                ctx_manager.reset(token)
                return
            except Exception as e:
                logger.warning(e, exc_info=True)
                ctx_manager.reset(token)
                raise

            # Reset BEFORE yielding to caller
            ctx_manager.reset(token)
            yield value

    return wrapper

WrapperStrategy

Bases: ABC

Abstract base strategy for method interception.

This interface implements the Strategy Pattern. It allows the instrumentor to decide what to do when a method is called, decoupling it from how to patch the method.

Source code in src/agent_cover/instrumentation/strategies.py 
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
class WrapperStrategy(ABC):
    """Abstract base strategy for method interception.

    This interface implements the **Strategy Pattern**. It allows the instrumentor
    to decide *what* to do when a method is called, decoupling it from *how* to
    patch the method.
    """

    @abstractmethod
    def wrap(self, original: Any, context: Any) -> Any:
        """Wraps a method or class.

        Args:
            original: The original method or class to be wrapped.
            context: A context object (can contain registry, context_manager, etc.).

        Returns:
            Any: The wrapped method or class.
        """
        pass
Functions
wrap(original, context) abstractmethod

Wraps a method or class.

Parameters:

Name Type Description Default
original Any

The original method or class to be wrapped.

 required
context Any

A context object (can contain registry, context_manager, etc.).

 required

Returns:

Name Type Description
Any Any

The wrapped method or class.
Source code in src/agent_cover/instrumentation/strategies.py 
30
31
32
33
34
35
36
37
38
39
40
41
@abstractmethod
def wrap(self, original: Any, context: Any) -> Any:
    """Wraps a method or class.

    Args:
        original: The original method or class to be wrapped.
        context: A context object (can contain registry, context_manager, etc.).

    Returns:
        Any: The wrapped method or class.
    """
    pass