fix(auth): isolate resolved credentials in context to prevent race conditions and data leakage

Co-authored-by: Xuan Yang <xygoogle@google.com>
PiperOrigin-RevId: 901957852

Author: xuanyang15

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/evaluation/eval_metrics.py

@@ -115,19 +115,6 @@ class BaseCriterion(BaseModel):
       description="The threshold to be used by the metric.",
   )
 
-  include_intermediate_responses_in_final: bool = Field(
-      default=False,
-      description=(
-          "Whether to evaluate the full agent response including intermediate"
-          " natural language text (e.g. text emitted before tool calls) in"
-          " addition to the final response. By default, only the final"
-          " response text is sent to the judge. When True, text from all"
-          " intermediate invocation events is concatenated with the final"
-          " response before evaluation. This is useful for agents that emit"
-          " text both before and after tool calls within a single invocation."
-      ),
-  )
-
 
 class LlmAsAJudgeCriterion(BaseCriterion):
   """Criterion when using LLM-As-A-Judge metric."""

src/google/adk/evaluation/llm_as_judge_utils.py

@@ -26,8 +26,6 @@
 from .common import EvalBaseModel
 from .eval_case import get_all_tool_calls_with_responses
 from .eval_case import IntermediateDataType
-from .eval_case import Invocation
-from .eval_case import InvocationEvents
 from .eval_metrics import RubricScore
 from .evaluator import EvalStatus
 
@@ -46,38 +44,8 @@ class Label(enum.Enum):
 
 
 def get_text_from_content(
-    content: Optional[Union[genai_types.Content, Invocation]],
-    *,
-    include_intermediate_responses_in_final: bool = False,
+    content: Optional[genai_types.Content],
 ) -> Optional[str]:
-  """Extracts text from a `Content` or an `Invocation`.
-
-  When `content` is a `Content`, returns the concatenated text of its parts.
-
-  When `content` is an `Invocation`, returns the text of the invocation's final
-  response. If `include_intermediate_responses_in_final` is True, text from
-  intermediate invocation events (e.g. natural language emitted before tool
-  calls) is concatenated with the final response text.
-  """
-  if isinstance(content, Invocation):
-    if not include_intermediate_responses_in_final:
-      # Flag off: revert to basic plain-Content behavior.
-      return get_text_from_content(content.final_response)
-
-    parts: list[str] = []
-    if isinstance(content.intermediate_data, InvocationEvents):
-      # Walk intermediate events in order; collect text parts.
-      for event in content.intermediate_data.invocation_events:
-        text = get_text_from_content(event.content)
-        if text:
-          parts.append(text)
-    # Then fetch the final response text and append it to the end.
-    final_text = get_text_from_content(content.final_response)
-    if final_text:
-      parts.append(final_text)
-
-    return "\n".join(parts) if parts else None
-
   if content and content.parts:
     return "\n".join([p.text for p in content.parts if p.text])
 

src/google/adk/evaluation/rubric_based_final_response_quality_v1.py

@@ -274,18 +274,7 @@ def format_auto_rater_prompt(
     """Returns the autorater prompt."""
     self.create_effective_rubrics_list(actual_invocation.rubrics)
     user_input = get_text_from_content(actual_invocation.user_content)
-
-    criterion = self._eval_metric.criterion
-    include_intermediate = getattr(
-        criterion, "include_intermediate_responses_in_final", False
-    )
-    final_response = (
-        get_text_from_content(
-            actual_invocation,
-            include_intermediate_responses_in_final=include_intermediate,
-        )
-        or ""
-    )
+    final_response = get_text_from_content(actual_invocation.final_response)
 
     rubrics_text = "\n".join([
         f"*  {r.rubric_content.text_property}"

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

@@ -143,10 +143,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 +159,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

src/google/adk/tools/mcp_tool/mcp_toolset.py

@@ -203,16 +203,31 @@ def __init__(
     )
     self._use_mcp_resources = use_mcp_resources
 
-  def _get_auth_headers(self) -> Optional[Dict[str, str]]:
+  def _get_auth_headers(
+      self, readonly_context: Optional[ReadonlyContext] = None
+  ) -> Optional[Dict[str, str]]:
     """Build authentication headers from exchanged credential.
 
+    Args:
+      readonly_context: Readonly context to get credentials from.
+
     Returns:
         Dictionary of auth headers, or None if no auth configured.
     """
-    if not self._auth_config or not self._auth_config.exchanged_auth_credential:
+    if not self._auth_config:
       return None
 
-    credential = self._auth_config.exchanged_auth_credential
+    credential = None
+    if readonly_context:
+      credential = readonly_context.get_credential(
+          self._auth_config.credential_key
+      )
+
+    if not credential:
+      credential = self._auth_config.exchanged_auth_credential
+
+    if not credential:
+      return None
     headers: Optional[Dict[str, str]] = None
 
     if credential.oauth2:
@@ -289,7 +304,7 @@ async def _execute_with_session(
         headers.update(provider_headers)
 
     # Add auth headers from exchanged credential if available
-    auth_headers = self._get_auth_headers()
+    auth_headers = self._get_auth_headers(readonly_context)
     if auth_headers:
       headers.update(auth_headers)
 

tests/unittests/auth/test_toolset_auth.py

@@ -110,6 +110,7 @@ def mock_invocation_context(self):
     ctx.credential_service = None
     ctx.app_name = "test-app"
     ctx.user_id = "test-user"
+    ctx.credential_by_key = {}
     return ctx
 
   @pytest.fixture
@@ -154,10 +155,10 @@ async def test_toolset_without_auth_config_skipped(
     assert mock_invocation_context.end_invocation is False
 
   @pytest.mark.asyncio
-  async def test_toolset_with_credential_available_populates_config(
+  async def test_toolset_with_credential_available_populates_context(
       self, mock_invocation_context, mock_agent
   ):
-    """Test that credential is populated in auth_config when available."""
+    """Test that credential is stored in invocation context when available."""
     auth_config = create_oauth2_auth_config()
     toolset = MockToolset(auth_config=auth_config)
     mock_agent.tools = [toolset]
@@ -184,8 +185,52 @@ async def test_toolset_with_credential_available_populates_config(
     # No auth request events - credential was available
     assert len(events) == 0
     assert mock_invocation_context.end_invocation is False
-    # Credential should be populated in auth_config
-    assert auth_config.exchanged_auth_credential == mock_credential
+    # Credential should be stored in invocation context, not auth_config
+    assert (
+        mock_invocation_context.credential_by_key[auth_config.credential_key]
+        == mock_credential
+    )
+    assert auth_config.exchanged_auth_credential is None
+
+  @pytest.mark.asyncio
+  async def test_toolset_auth_uses_copy_and_does_not_mutate_shared_config(
+      self, mock_invocation_context, mock_agent
+  ):
+    """Test that _resolve_toolset_auth uses a copy and does not mutate shared config."""
+    auth_config = create_oauth2_auth_config()
+    toolset = MockToolset(auth_config=auth_config)
+    mock_agent.tools = [toolset]
+
+    def create_mock_cm(cfg):
+      m = AsyncMock()
+      m._auth_config = cfg
+
+      async def get_cred(ctx):
+        cfg.exchanged_auth_credential = AuthCredential(
+            auth_type=AuthCredentialTypes.OAUTH2,
+            oauth2=OAuth2Auth(auth_uri="https://example.com/consent"),
+        )
+        return None
+
+      m.get_auth_credential = AsyncMock(side_effect=get_cred)
+      return m
+
+    with patch(
+        "google.adk.flows.llm_flows.base_llm_flow.CredentialManager",
+        side_effect=create_mock_cm,
+    ):
+      events = []
+      async for event in _resolve_toolset_auth(
+          mock_invocation_context, mock_agent
+      ):
+        events.append(event)
+
+    # Should yield one auth request event
+    assert len(events) == 1
+    assert mock_invocation_context.end_invocation is True
+
+    # The shared auth_config should NOT be mutated
+    assert auth_config.exchanged_auth_credential is None
 
   @pytest.mark.asyncio
   async def test_toolset_without_credential_yields_auth_event(

tests/unittests/evaluation/test_llm_as_judge_utils.py

@@ -19,7 +19,6 @@
 from google.adk.evaluation.app_details import AgentDetails
 from google.adk.evaluation.app_details import AppDetails
 from google.adk.evaluation.eval_case import IntermediateData
-from google.adk.evaluation.eval_case import Invocation
 from google.adk.evaluation.eval_case import InvocationEvent
 from google.adk.evaluation.eval_case import InvocationEvents
 from google.adk.evaluation.eval_rubrics import RubricScore
@@ -89,49 +88,6 @@ def test_get_text_from_content_with_mixed_parts():
   assert get_text_from_content(content) == "Hello\nWorld"
 
 
-def test_get_text_from_content_with_invocation_include_intermediate_responses_in_final():
-  """Tests get_text_from_content on an Invocation with and without the flag."""
-  intermediate_text = "Let me check."
-  final_response_text = "Done."
-  invocation = Invocation(
-      user_content=genai_types.Content(parts=[genai_types.Part(text="user")]),
-      intermediate_data=InvocationEvents(
-          invocation_events=[
-              InvocationEvent(
-                  author="agent",
-                  content=genai_types.Content(
-                      parts=[genai_types.Part(text=intermediate_text)]
-                  ),
-              ),
-              InvocationEvent(
-                  author="tool",
-                  content=genai_types.Content(
-                      parts=[
-                          genai_types.Part(
-                              function_call=genai_types.FunctionCall(name="t")
-                          )
-                      ]
-                  ),
-              ),
-          ]
-      ),
-      final_response=genai_types.Content(
-          parts=[genai_types.Part(text=final_response_text)]
-      ),
-  )
-
-  # Flag off (default): only the final response text is returned.
-  assert get_text_from_content(invocation) == final_response_text
-
-  # Flag on: intermediate text is concatenated before the final response.
-  assert (
-      get_text_from_content(
-          invocation, include_intermediate_responses_in_final=True
-      )
-      == f"{intermediate_text}\n{final_response_text}"
-  )
-
-
 def test_get_eval_status_with_none_score():
   """Tests get_eval_status returns NOT_EVALUATED for a None score."""
   assert get_eval_status(score=None, threshold=0.5) == EvalStatus.NOT_EVALUATED

tests/unittests/tools/mcp_tool/test_mcp_toolset_auth.py

@@ -267,3 +267,32 @@ def test_get_auth_headers_api_key_non_header_logs_warning(self, caplog):
 
     # Should return None for non-header API key
     assert headers is None
+
+  def test_get_auth_headers_reads_from_readonly_context(
+      self, toolset_with_oauth2
+  ):
+    """Test that _get_auth_headers reads from ReadonlyContext first."""
+    from google.adk.agents.readonly_context import ReadonlyContext
+
+    mock_readonly_context = Mock(spec=ReadonlyContext)
+    mock_credential = AuthCredential(
+        auth_type=AuthCredentialTypes.OAUTH2,
+        oauth2=OAuth2Auth(access_token="token-from-context"),
+    )
+    mock_readonly_context.get_credential.return_value = mock_credential
+
+    # Even if exchanged_auth_credential has a different value
+    toolset_with_oauth2._auth_config.exchanged_auth_credential = AuthCredential(
+        auth_type=AuthCredentialTypes.OAUTH2,
+        oauth2=OAuth2Auth(access_token="token-from-config"),
+    )
+
+    headers = toolset_with_oauth2._get_auth_headers(
+        readonly_context=mock_readonly_context
+    )
+
+    assert headers is not None
+    assert headers["Authorization"] == "Bearer token-from-context"
+    mock_readonly_context.get_credential.assert_called_once_with(
+        toolset_with_oauth2._auth_config.credential_key
+    )