should InvocationContext.end_invocation skip the after_agent_callbacks?

[rcleveng]

Right now, if a tool_call sets end_invocation to true, then the agent's after_agent_callbacks will never be called (in base_agent.py it has the following code:

with tracer.start_as_current_span(f'agent_run [{self.name}]'):
      ctx = self._create_invocation_context(parent_context)

      if event := await self.__handle_before_agent_callback(ctx):
        yield event
      if ctx.end_invocation:
        return

      async for event in self._run_async_impl(ctx):
        yield event

      if ctx.end_invocation:
        return

      if event := await self.__handle_after_agent_callback(ctx):
        yield event

I'm using end_invocation in a tool so that the results of the tool call get returned verbatim and doesn't loop through the LLM again.

This is surprising to me, but maybe it's expected. I think this is breaking the Opik integration for tracing.

[akshatshaw]

can you show the code example where you have implemented this, i am also facing the comet issue.

[rcleveng]

Sure, here's the entirety of the change:

Sent as a PR: #1761

diff --git a/src/google/adk/agents/base_agent.py b/src/google/adk/agents/base_agent.py
index bdc10ac..c4ce18e 100644
--- a/src/google/adk/agents/base_agent.py
+++ b/src/google/adk/agents/base_agent.py
@@ -147,8 +147,9 @@ class BaseAgent(BaseModel):
       async for event in self._run_async_impl(ctx):
         yield event
 
-      if ctx.end_invocation:
-        return
+      # Don't end invocation here so the afer agent callbacks are still called
+      # if ctx.end_invocation:
+      #   return
 
       if event := await self.__handle_after_agent_callback(ctx):
         yield event