Merge branch 'main' into fix/databasesession-update-wrong-timezone

Author: rohityan

README.md

@@ -69,6 +69,11 @@ You can install the latest stable version of ADK using `pip`:
 pip install google-adk
 ```
 
+To install optional integrations, you can use the following command:
+```bash
+pip install "google-adk[extensions]"
+```
+
 The release cadence is roughly bi-weekly.
 
 This version is recommended for most users as it represents the most recent official release.

contributing/samples/adk_triaging_agent/agent.py

@@ -37,7 +37,7 @@
     "models": "xuanyang15",
     "services": "DeanChensj",
     "tools": "xuanyang15",
-    "tracing": "jawoszek",
+    "tracing": "mhenc",
     "web": "wyf7107",
     "workflow": "DeanChensj",
 }

src/google/adk/agents/base_agent.py

@@ -16,6 +16,7 @@
 
 import inspect
 import logging
+import sys
 from typing import Any
 from typing import AsyncGenerator
 from typing import Awaitable
@@ -285,7 +286,9 @@ async def run_async(
       Event: the events generated by the agent.
     """
 
-    with tracer.start_as_current_span(f'invoke_agent {self.name}') as span:
+    cm = tracer.start_as_current_span(f'invoke_agent {self.name}')
+    span = cm.__enter__()
+    try:
       ctx = self._create_invocation_context(parent_context)
       tracing.trace_agent_invocation(span, self, ctx)
       if event := await self._handle_before_agent_callback(ctx):
@@ -302,6 +305,23 @@ 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/agents/invocation_context.py

@@ -28,6 +28,7 @@
 from ..apps.app import EventsCompactionConfig
 from ..apps.app import ResumabilityConfig
 from ..artifacts.base_artifact_service import BaseArtifactService
+from ..auth.auth_credential import AuthCredential
 from ..auth.credential_service.base_credential_service import BaseCredentialService
 from ..events.event import Event
 from ..memory.base_memory_service import BaseMemoryService
@@ -214,6 +215,9 @@ class InvocationContext(BaseModel):
   canonical_tools_cache: Optional[list[BaseTool]] = None
   """The cache of canonical tools for this invocation."""
 
+  credential_by_key: dict[str, AuthCredential] = Field(default_factory=dict)
+  """The resolved credentials for this invocation, keyed by credential_key."""
+
   _invocation_cost_manager: _InvocationCostManager = PrivateAttr(
       default_factory=_InvocationCostManager
   )

src/google/adk/agents/readonly_context.py

@@ -22,6 +22,7 @@
 if TYPE_CHECKING:
   from google.genai import types
 
+  from ..auth.auth_credential import AuthCredential
   from ..sessions.session import Session
   from .invocation_context import InvocationContext
   from .run_config import RunConfig
@@ -69,3 +70,7 @@ def user_id(self) -> str:
   def run_config(self) -> Optional[RunConfig]:
     """The run config of the current invocation. READONLY field."""
     return self._invocation_context.run_config
+
+  def get_credential(self, key: str) -> Optional[AuthCredential]:
+    """Gets a resolved credential by key for this invocation."""
+    return self._invocation_context.credential_by_key.get(key)

src/google/adk/cli/adk_web_server.py

@@ -2075,6 +2075,7 @@ async def run_agent_live(
         proactive_audio: bool | None = Query(default=None),
         enable_affective_dialog: bool | None = Query(default=None),
         enable_session_resumption: bool | None = Query(default=None),
+        save_live_blob: bool = Query(default=False),
     ) -> None:
       ws_origin = websocket.headers.get("origin")
       if ws_origin is not None and not _is_request_origin_allowed(
@@ -2117,6 +2118,7 @@ async def forward_events():
                 if enable_session_resumption is not None
                 else None
             ),
+            save_live_blob=save_live_blob,
         )
         async with Aclosing(
             runner.run_live(

src/google/adk/cli/browser/index.html

@@ -123,23 +123,24 @@ async def execute(
     )
 
   @override
-  async def read_file(self, path: str) -> bytes:
+  async def read_file(self, path: str | Path) -> bytes:
     if self._working_dir is None:
       raise RuntimeError('`working_dir` is not set. Call initialize() first.')
 
-    path = self._resolve_path(path)
-    return await asyncio.to_thread(self._sync_read, path)
+    resolved = self._resolve_path(path)
+    return await asyncio.to_thread(self._sync_read, resolved)
 
   @override
-  async def write_file(self, path: str, content: str | bytes) -> None:
+  async def write_file(self, path: str | Path, content: str | bytes) -> None:
     if self._working_dir is None:
       raise RuntimeError('`working_dir` is not set. Call initialize() first.')
 
-    path = self._resolve_path(path)
-    return await asyncio.to_thread(self._sync_write, path, content)
+    resolved = self._resolve_path(path)
+    return await asyncio.to_thread(self._sync_write, resolved, content)
 
-  def _resolve_path(self, path: str) -> str:
+  def _resolve_path(self, path: str | Path) -> str:
     """Resolve a relative path against the working directory."""
+    path = str(path)
     if os.path.isabs(path):
       return path
     return os.path.join(self._working_dir, path)

…/google/adk/cli/browser/main-3OBAHHYV.js …/google/adk/cli/browser/main-TCIQIOZ3.jssrc/google/adk/cli/browser/main-3OBAHHYV.js renamed to src/google/adk/cli/browser/main-TCIQIOZ3.js src/google/adk/cli/browser/main-3OBAHHYV.js renamed to src/google/adk/cli/browser/main-TCIQIOZ3.js

@@ -18,6 +18,7 @@
 import asyncio
 import inspect
 import logging
+import sys
 from typing import AsyncGenerator
 from typing import Optional
 from typing import TYPE_CHECKING
@@ -143,10 +144,11 @@ async def _resolve_toolset_auth(
     if not auth_config:
       continue
 
+    auth_config_copy = auth_config.model_copy(deep=True)
     try:
-      credential = await CredentialManager(auth_config).get_auth_credential(
-          callback_context
-      )
+      credential = await CredentialManager(
+          auth_config_copy
+      ).get_auth_credential(callback_context)
     except ValueError as e:
       # Validation errors from CredentialManager should be logged but not
       # block the flow - the toolset may still work without auth
@@ -158,14 +160,16 @@ async def _resolve_toolset_auth(
       credential = None
 
     if credential:
-      # Populate in-place for toolset to use in get_tools()
-      auth_config.exchanged_auth_credential = credential
+      # Store in invocation context to avoid data leakage and race conditions
+      invocation_context.credential_by_key[auth_config.credential_key] = (
+          credential
+      )
     else:
       # Need auth - will interrupt
       toolset_id = (
           f'{TOOLSET_AUTH_CREDENTIAL_ID_PREFIX}{type(tool_union).__name__}'
       )
-      pending_auth_requests[toolset_id] = auth_config
+      pending_auth_requests[toolset_id] = auth_config_copy
 
   if not pending_auth_requests:
     return
@@ -1165,7 +1169,9 @@ async def _call_llm_async(
   ) -> AsyncGenerator[LlmResponse, None]:
 
     async def _call_llm_with_tracing() -> AsyncGenerator[LlmResponse, None]:
-      with tracer.start_as_current_span('call_llm') as span:
+      cm = tracer.start_as_current_span('call_llm')
+      span = cm.__enter__()
+      try:
         # 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(
@@ -1258,6 +1264,23 @@ 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: