fix: revert fix handle ContextVar detach errors during task cancellation

Co-authored-by: Max Ind <maxind@google.com>
PiperOrigin-RevId: 903267607

Author: RKest

src/google/adk/agents/base_agent.py

@@ -16,7 +16,6 @@
 
 import inspect
 import logging
-import sys
 from typing import Any
 from typing import AsyncGenerator
 from typing import Awaitable
@@ -286,9 +285,7 @@ async def run_async(
       Event: the events generated by the agent.
     """
 
-    cm = tracer.start_as_current_span(f'invoke_agent {self.name}')
-    span = cm.__enter__()
-    try:
+    with tracer.start_as_current_span(f'invoke_agent {self.name}') as span:
       ctx = self._create_invocation_context(parent_context)
       tracing.trace_agent_invocation(span, self, ctx)
       if event := await self._handle_before_agent_callback(ctx):
@@ -305,23 +302,6 @@ async def run_async(
 
       if event := await self._handle_after_agent_callback(ctx):
         yield event
-    except BaseException:
-      try:
-        cm.__exit__(*sys.exc_info())
-      except ValueError:
-        logger.warning(
-            'Failed to detach context during generator cleanup, likely due to'
-            ' cancellation.'
-        )
-      raise
-    else:
-      try:
-        cm.__exit__(None, None, None)
-      except ValueError:
-        logger.warning(
-            'Failed to detach context during generator cleanup, likely due to'
-            ' cancellation.'
-        )
 
   @final
   async def run_live(

src/google/adk/flows/llm_flows/base_llm_flow.py

@@ -18,7 +18,6 @@
 import asyncio
 import inspect
 import logging
-import sys
 from typing import AsyncGenerator
 from typing import Optional
 from typing import TYPE_CHECKING
@@ -1169,9 +1168,7 @@ async def _call_llm_async(
   ) -> AsyncGenerator[LlmResponse, None]:
 
     async def _call_llm_with_tracing() -> AsyncGenerator[LlmResponse, None]:
-      cm = tracer.start_as_current_span('call_llm')
-      span = cm.__enter__()
-      try:
+      with tracer.start_as_current_span('call_llm') as span:
         # Runs before_model_callback inside the call_llm span so
         # plugins observe the same span as after/error callbacks.
         if response := await self._handle_before_model_callback(
@@ -1264,23 +1261,6 @@ async def _call_llm_with_tracing() -> AsyncGenerator[LlmResponse, None]:
                   llm_response = altered
 
               yield llm_response
-      except BaseException:
-        try:
-          cm.__exit__(*sys.exc_info())
-        except ValueError:
-          logger.warning(
-              'Failed to detach context during generator cleanup, likely due to'
-              ' cancellation.'
-          )
-        raise
-      else:
-        try:
-          cm.__exit__(None, None, None)
-        except ValueError:
-          logger.warning(
-              'Failed to detach context during generator cleanup, likely due to'
-              ' cancellation.'
-          )
 
     async with Aclosing(_call_llm_with_tracing()) as agen:
       async for event in agen:

src/google/adk/runners.py

@@ -543,9 +543,7 @@ async def _run_with_trace(
         new_message: Optional[types.Content] = None,
         invocation_id: Optional[str] = None,
     ) -> AsyncGenerator[Event, None]:
-      cm = tracer.start_as_current_span('invocation')
-      span = cm.__enter__()
-      try:
+      with tracer.start_as_current_span('invocation'):
         session = await self._get_or_create_session(
             user_id=user_id,
             session_id=session_id,
@@ -629,23 +627,6 @@ async def execute(ctx: InvocationContext) -> AsyncGenerator[Event]:
               self.session_service,
               skip_token_compaction=invocation_context.token_compaction_checked,
           )
-      except BaseException:
-        try:
-          cm.__exit__(*sys.exc_info())
-        except ValueError:
-          logger.warning(
-              'Failed to detach context during generator cleanup, likely due to'
-              ' cancellation.'
-          )
-        raise
-      else:
-        try:
-          cm.__exit__(None, None, None)
-        except ValueError:
-          logger.warning(
-              'Failed to detach context during generator cleanup, likely due to'
-              ' cancellation.'
-          )
 
     async with Aclosing(_run_with_trace(new_message, invocation_id)) as agen:
       async for event in agen:

tests/unittests/tools/BUILD

@@ -31,17 +31,3 @@ pytest_test(
         "//third_party/py/pytest_asyncio",
     ],
 )
-
-pytest_test(
-    name = "test_context_cancellation",
-    srcs = ["test_context_cancellation.py"],
-    args = [
-        "-p",
-        "pytest_asyncio.plugin",
-    ],
-    deps = [
-        "//third_party/py/google/adk",
-        "//third_party/py/google/genai",
-        "//third_party/py/pytest_asyncio",
-    ],
-)