Merge remote-tracking branch 'origin/main' into fix-thought-signature-pruning

Author: its-amann

src/google/adk/integrations/agent_identity/_iam_connector_credentials_provider.py

@@ -238,7 +238,7 @@ async def get_auth_credential(
       logger.debug("Auth credential obtained immediately.")
       return _construct_auth_credential(response)
 
-    if metadata and metadata.consent_pending:
+    if metadata is not None and "consent_pending" in metadata:
       # Get 2-legged OAuth token. Allow enough time for token exchange.
       try:
         operation = await self._poll_credentials(

src/google/adk/telemetry/google_cloud.py

@@ -119,7 +119,6 @@ def get_gcp_exporters(
   if enable_cloud_logging:
     exporter = _get_gcp_logs_exporter(
         project_id=project_id,
-        credentials=credentials,
     )
     if exporter:
       log_record_processors.append(exporter)
@@ -187,11 +186,9 @@ def _get_gcp_metrics_exporter(project_id: str) -> MetricReader:
 
 def _get_gcp_logs_exporter(
     project_id: str,
-    credentials: Optional["Credentials"] = None,
 ) -> LogRecordProcessor:
   if os.getenv("GOOGLE_CLOUD_AGENT_ENGINE_ID"):
     return _get_agent_engine_logs_exporter(
-        credentials=credentials,
         project_id=project_id,
     )
 
@@ -338,18 +335,14 @@ def _use_client_cert_effective() -> bool:
 
 def _get_agent_engine_logs_exporter(
     *,
-    credentials: "Credentials",
     project_id: str,
 ):
   """Configures logging for Agent Engine.
 
   Args:
-    credentials: Credentials to use for export calls.
     project_id: Project to which to write logs.
   """
   try:
-    from google.cloud.logging_v2.services import logging_service_v2
-    from google.cloud.logging_v2.services.logging_service_v2.transports import grpc
     from opentelemetry.exporter import cloud_logging
   except (ImportError, AttributeError):
     logger.warning(
@@ -363,46 +356,21 @@ def _get_agent_engine_logs_exporter(
     )
     return
 
-  if "gen_ai_latest_experimental" in os.getenv(
-      "OTEL_SEMCONV_STABILITY_OPT_IN", ""
-  ).split(","):
-    # Specify credentials to avoid expensive call to `google.auth.default()`
-    channel = grpc.LoggingServiceV2GrpcTransport.create_channel(
-        credentials=credentials,
-        # pylint: disable-next=protected-access
-        options=cloud_logging._OPTIONS,
-    )
-    return BatchLogRecordProcessor(
-        cloud_logging.CloudLoggingExporter(
-            client=logging_service_v2.LoggingServiceV2Client(
-                transport=grpc.LoggingServiceV2GrpcTransport(
-                    credentials=credentials,
-                    channel=channel,
-                ),
-            ),
-            project_id=project_id,
-            default_log_name=os.getenv(
-                "GCP_DEFAULT_LOG_NAME", "adk-on-agent-engine"
-            ),
-        ),
-    )
-  else:
-
-    class _SimpleLogRecordProcessor(SimpleLogRecordProcessor):
-
-      def force_flush(
-          self, timeout_millis: int = 30000
-      ) -> bool:  # pylint: disable=no-self-use
-        _ = sys.stdout.flush()
-        _ = sys.stderr.flush()
-        return super().force_flush()
-
-    return _SimpleLogRecordProcessor(
-        cloud_logging.CloudLoggingExporter(
-            project_id=project_id,
-            default_log_name=os.getenv(
-                "GCP_DEFAULT_LOG_NAME", "adk-on-agent-engine"
-            ),
-            structured_json_file=sys.stdout,
-        ),
-    )
+  class _SimpleLogRecordProcessor(SimpleLogRecordProcessor):
+
+    def force_flush(
+        self, timeout_millis: int = 30000
+    ) -> bool:  # pylint: disable=no-self-use
+      _ = sys.stdout.flush()
+      _ = sys.stderr.flush()
+      return super().force_flush()
+
+  return _SimpleLogRecordProcessor(
+      cloud_logging.CloudLoggingExporter(
+          project_id=project_id,
+          default_log_name=os.getenv(
+              "GCP_DEFAULT_LOG_NAME", "adk-on-agent-engine"
+          ),
+          structured_json_file=sys.stdout,
+      ),
+  )

tests/unittests/integrations/agent_identity/test_iam_connector_credentials_provider.py

@@ -469,3 +469,42 @@ async def test_get_auth_credential_raises_error_if_consent_canceled(
       RuntimeError, match="Failed to retrieve consent based credential."
   ):
     await provider.get_auth_credential(auth_scheme, context)
+
+
+async def test_get_auth_credential_handles_consent_pending_state_correctly(
+    mock_client,
+    auth_scheme,
+    context,
+    provider,
+):
+  """Test get_auth_credential enters the polling loop when consent is pending."""
+  # 1. Setup the first retrieve_credentials call to return a pending Operation
+  # containing consent_pending metadata.
+  meta_pb = RetrieveCredentialsMetadata.pb()()
+  meta_pb.consent_pending.SetInParent()
+
+  op_pending = Operation(done=False)
+  op_pending.metadata.value = meta_pb.SerializeToString()
+
+  # 2. Setup the second retrieve_credentials call (the poll) to return success.
+  op_success = Operation(done=True)
+  resp = RetrieveCredentialsResponse(
+      header="Authorization: Bearer", token="valid-token"
+  )
+  op_success.response.value = RetrieveCredentialsResponse.serialize(resp)
+
+  # Configure mock client to return pending first, then success
+  mock_client.retrieve_credentials.side_effect = [
+      Mock(operation=op_pending),
+      Mock(operation=op_success),
+  ]
+
+  # 3. Call the provider
+  credential = await provider.get_auth_credential(auth_scheme, context)
+
+  # 4. Verify that it polled and successfully returned the token
+  assert credential is not None
+  assert credential.http.credentials.token == "valid-token"
+
+  # Verify that retrieve_credentials was called twice (initial + 1 poll)
+  assert mock_client.retrieve_credentials.call_count == 2

tests/unittests/telemetry/test_google_cloud.py

@@ -62,7 +62,7 @@ def test_get_gcp_exporters(
   )
   monkeypatch.setattr(
       "google.adk.telemetry.google_cloud._get_gcp_logs_exporter",
-      lambda project_id, credentials: mock.MagicMock(),
+      lambda project_id: mock.MagicMock(),
   )
 
   # Act.