Python: Is wrapping an executor's handler maps a safe way to run code around every handler call

[Banka Uday (banka-uday-4129)]

Category: Q&A

Title: Python: Is wrapping an executor's handler maps a safe way to run code around every handler call?

What we're trying to do

We run workflows built from our own executors plus AgentExecutor subclasses. Each time an executor handles a message, we want to run our own code around that handler call, so that we can:

• record the execution (start, end, error, and whether it was a normal message or the response to a request_info) in our own store;
• make that execution visible to everything running inside it. We set a ContextVar, so that the agent's chat and function middleware can link each model call and tool call to the executor execution that made it;
• add the executor's input and output to its executor.process span.

There are agent, chat and function middleware, but we couldn't find a hook at the executor level, so we'd like to check that our approach is sound.

What we plan to do

In a base class shared by our executors, once super().__init__() has discovered the handlers, wrap every value of _handlers and _response_handlers. Simplified:

current_execution: ContextVar[Any] = ContextVar("current_execution", default=None)


def around_handler(handler, *, executor_id: str, message_type: str):
    @functools.wraps(handler)
    async def wrapped(*args, **kwargs):
        execution = await before(executor_id, message_type)   # our code
        token = current_execution.set(execution)           # read by our chat/function middleware
        try:
            await handler(*args, **kwargs)
        except BaseException as error:
            await on_error(execution, error)                # our code
            raise
        else:
            await after(execution)                          # our code
        finally:
            current_execution.reset(token)
    return wrapped


class MyExecutorBase(Executor):
    def wrap_handlers(self) -> None:
        for key, handler in self._handlers.items():
            self._handlers[key] = around_handler(handler, executor_id=self.id, message_type="standard")
        for key, handler in self._response_handlers.items():
            self._response_handlers[key] = around_handler(handler, executor_id=self.id, message_type="response")

Why we chose this, from reading agent-framework-core 1.24.0:

• Executor.__init__ stores the discovered handlers in self._handlers (_workflows/_executor.py ~L221) and self._response_handlers (_request_info_mixin.py ~L81).
• Executor.execute() takes the execution lock, opens the executor.process span, looks the handler up in those maps (_find_handler / _find_response_handler, which binds the original request with functools.partial), and calls await handler(message, context) (_executor.py ~L276–330).
• So the wrapper runs once per handler call, in the same task as the handler, inside the lock and the span. It also covers inherited handlers such as AgentExecutor.handle_user_input_response, and routing is unchanged because the dict keys stay the same.
• A handler that calls another handler directly (for example self.run(...)) doesn't go through the maps, so it isn't counted twice.

Other options we looked at

• Workflow events (executor_invoked / executor_completed / executor_failed): they're read in the task that consumes the run stream, not the one running the handler. A ContextVar set there isn't visible to the agent's chat or function middleware, and the event may arrive after the model calls it should be linked to.
• AgentMiddleware: only covers agent executors, not our other executors. It also misses failures outside agent.run and response-handler calls that don't reach the agent.
• A SpanProcessor on executor.process: only works while instrumentation is on and nothing is sampled out. Its callbacks are synchronous, so we can't await our store there.
• Overriding Executor.execute(): its docstring says not to. It would also run before the lock and the span.
• Decorating each handler: works, but every handler, including the ones inherited from the framework, would have to be redeclared, and the decorator must keep the signature used for type-based routing.

Questions

1. Is wrapping the handler maps like this safe to rely on across 1.x releases?
2. Could it break anything we haven't considered: checkpointing and restore, handler discovery, the data on executor_completed, edge or fan-in delivery, or anything else in execute()?
3. Is there a better or supported way to run code around every handler call of an executor that we've missed?

[Vishal Punjabi (vshalpnjabi)]

Not a maintainer, but I read through _workflows/_executor.py on main and your reading of execute() matches what I see. A few things to watch for with the wrap-after-__init__ approach:

1. defer_discovery=True. FunctionExecutor (_function_executor.py L148) calls super().__init__(..., defer_discovery=True), then registers its handler with _register_instance_handler and calls _discover_handlers() / _discover_response_handlers() itself (L173-184). Some internal executors in _edge.py (L260, L285) do the same. If any of your executors take that path, or anything calls _discover_* again later, the maps get refilled with unwrapped methods. Wrapping lazily (on first execute, or keyed off a flag) is safer than doing it once after super().__init__().

2. Event ordering. executor_invoked is added before await handler(...) and executor_completed/executor_failed after it, so your before/after run between those events, not around them. That's fine for linking, just don't expect your record to bracket the events.

3. BaseException vs Exception. execute() only emits executor_failed for Exception. Your wrapper also calls on_error on CancelledError, which you'll see when a run is cancelled. Probably what you want, but your store will have "errors" the workflow never reports as failures.

4. Checkpointing. Restore goes through on_checkpoint_save/on_checkpoint_restore (state dicts) and doesn't rebuild the handler maps, so wrapping doesn't affect it. Routing and _handler_specs are untouched since you keep the same keys.

5. Response handlers. _find_response_handler returns functools.partial(handler, request) around whatever is in the map, so your wrapper receives (request, response, ctx). That works with *args, just keep it in mind if before() wants to inspect the message.

On stability: _handlers and _response_handlers are private, so nothing guarantees them across 1.x. I'd pin a test that builds one of your executors and asserts every value in both maps is wrapped, so a refactor fails loudly instead of quietly dropping your records. If you want a supported hook, an issue asking for an executor-level middleware is probably the way to get it.

[Gianni Luigi Messina (Gianni300x)]

Two additional Python-level edge cases to consider alongside the private-API caveat above:

1. Child task lifetime. An asyncio.create_task() created inside the handler normally receives a copy of the current context. Resetting the token in the parent does not clear that copy, so a background task that outlives the handler can continue to see the old execution. If that is undesirable, make child work finish before the handler returns, or create the task with an explicitly chosen context. See the asyncio task documentation.
2. Failures in the instrumentation itself. If on_error() raises, it replaces the handler's original exception; if after() raises, an otherwise successful handler appears to fail. Decide explicitly whether persistence failures should fail the workflow, and test that policy for ordinary errors and cancellation. In particular, avoid accidentally suppressing cancellation while attempting async cleanup. If the wrapper is intended to be transparent, retain and return the handler's result (result = await handler(...); ...; return result) rather than discarding it.

These are observations about the simplified wrapper, not a claim that the private maps are a supported extension point. I have not run your workflow or tested checkpoint/restore behavior.