Merge branch 'main' into fix/dev-server-eval-set-legacy-nameerror

Author: gaurav-gandhi-2411

src/google/adk/plugins/auto_tracing_helpers.py

@@ -17,6 +17,7 @@
 from __future__ import annotations
 
 import asyncio
+from collections.abc import Mapping
 import dataclasses
 import functools
 import inspect
@@ -41,6 +42,77 @@
 _SCALAR_TYPES = frozenset({int, float, bool, str, bytes, type(None)})
 _DEFAULT_REPR_RE = re.compile(r"^<.+ object at 0x[0-9a-fA-F]+>$")
 
+# Types whose repr() renders live secrets (tokens, keys, passwords). Matched by
+# name over the MRO so this module never imports ``google.adk.auth``.
+_CREDENTIAL_TYPE_NAMES = frozenset({
+    "AuthConfig",
+    "AuthCredential",
+    "AuthToolArguments",
+    "Credentials",
+    "HttpAuth",
+    "HttpCredentials",
+    "OAuth2Auth",
+    "OAuth2Session",
+    "ServiceAccount",
+    "ServiceAccountCredential",
+})
+# Parameter names that conventionally carry secret material.
+_CREDENTIAL_ARG_NAMES = frozenset({
+    "api_key",
+    "auth_config",
+    "auth_credential",
+    "authorization",
+    "cookie",
+    "cookies",
+    "credential",
+    "credentials",
+    "password",
+    "private_key",
+    "secret",
+    "token",
+})
+_CREDENTIAL_ARG_SUFFIXES = (
+    "_api_key",
+    "_auth_config",
+    "_authorization",
+    "_cookie",
+    "_cookies",
+    "_credential",
+    "_credentials",
+    "_password",
+    "_private_key",
+    "_secret",
+    "_token",
+)
+# Bounds for the structural walk below. Both are deliberately generous: only
+# containers and objects consume node budget, so a list of a million ints
+# costs one node.
+_MAX_REDACT_DEPTH = 10
+_MAX_REDACT_NODES = 1024
+
+
+def _mro_holds_credential(cls: type) -> bool:
+  """True iff ``cls`` or one of its bases is a credential-bearing type."""
+  return any(k.__name__ in _CREDENTIAL_TYPE_NAMES for k in cls.__mro__)
+
+
+# Cached because the walk asks this of every non-scalar node it visits. The
+# annotation is spelled out because lru_cache erases the wrapped signature to
+# ``*args: Hashable``, which the ``type(value)`` the callers pass does not
+# satisfy.
+_is_credential_type: Callable[[type], bool] = functools.lru_cache(maxsize=512)(
+    _mro_holds_credential
+)
+
+
+@functools.lru_cache(maxsize=1024)
+def _is_credential_arg_name(name: str) -> bool:
+  """True iff a parameter called ``name`` conventionally holds a secret."""
+  lowered = name.lower()
+  return lowered in _CREDENTIAL_ARG_NAMES or lowered.endswith(
+      _CREDENTIAL_ARG_SUFFIXES
+  )
+
 
 @dataclasses.dataclass(frozen=True)
 class Caps:
@@ -73,8 +145,158 @@ def __repr__(self) -> str:
     )
 
 
+def _plain_repr(value: Any) -> str:
+  """``repr(value)`` that never raises."""
+  try:
+    return repr(value)
+  except Exception:  # pylint: disable=broad-exception-caught
+    return f"<unrepr-able {type(value).__name__}>"
+
+
+def _redacted_repr(value: Any) -> str | None:
+  """Renders ``value`` with nested credentials masked, or ``None`` if clean.
+
+  ``None`` means "no secret material anywhere in here", and the caller keeps
+  plain ``repr()``. Otherwise the walk rebuilds the rendering element by
+  element -- through mappings, sequences, sets, NamedTuples, dataclasses,
+  pydantic models and plain objects -- so a credential is masked wherever it
+  sits rather than only at the top level. Clean subtrees are still rendered
+  with ``repr()``, so the text of an ordinary value is unchanged.
+
+  The walk is bounded three ways: nesting depth, the number of container
+  nodes visited, and an id set that stops cycles. A subtree it refuses to
+  walk is elided rather than repr'd, so hitting a bound can never uncover a
+  secret. An object that hides state behind a leading underscore is
+  *inspected* there but only ever *rendered* from its public attributes, so
+  the redacted form never shows more than the original repr would have.
+  """
+  budget = [_MAX_REDACT_NODES]
+  active: set[int] = set()
+
+  def member(name: Any, v: Any, depth: int) -> str | None:
+    """Like ``walk`` but masks by field/key name too."""
+    if isinstance(name, str) and _is_credential_arg_name(name):
+      return f"<{type(v).__name__}>"
+    return walk(v, depth)
+
+  def members(items: Any, depth: int) -> list[str] | None:
+    """``["name=text", ...]``, or ``None`` when nothing needed masking.
+
+    Clean children are only rendered once the node is known to be dirty, so a
+    value with no secret in it costs a traversal and not a repr per node.
+    """
+    walked = [(name, v, member(name, v, depth)) for name, v in items]
+    if all(text is None for _, _, text in walked):
+      return None
+    return [
+        f"{name}={text if text is not None else _plain_repr(v)}"
+        for name, v, text in walked
+    ]
+
+  def walk(v: Any, depth: int) -> str | None:
+    if type(v) in _SCALAR_TYPES:
+      return None
+    cls = type(v)
+    if _is_credential_type(cls):
+      return f"<{cls.__name__}>"
+    if isinstance(v, type) or inspect.ismodule(v):
+      return None
+    # StreamResult already renders each sampled item through safe_repr, so
+    # walking it would only cost the yield count its own repr reports.
+    if isinstance(v, StreamResult):
+      return None
+    budget[0] -= 1
+    marker = id(v)
+    if budget[0] < 0 or depth >= _MAX_REDACT_DEPTH or marker in active:
+      return f"<{cls.__name__} ...>"
+    active.add(marker)
+    try:
+      return descend(v, depth + 1)
+    finally:
+      active.discard(marker)
+
+  def descend(v: Any, depth: int) -> str | None:
+    cls = type(v)
+    name = cls.__name__
+    speaks_for_itself = getattr(cls, "__repr__", None) is not object.__repr__
+    if speaks_for_itself and isinstance(v, tuple) and hasattr(cls, "_fields"):
+      parts = members(zip(cls._fields, v), depth)
+      return f"{name}({', '.join(parts)})" if parts else None
+    if speaks_for_itself and isinstance(v, Mapping):
+      walked = [
+          (k, walk(k, depth), item, member(k, item, depth))
+          for k, item in v.items()
+      ]
+      if all(kt is None and vt is None for _, kt, _, vt in walked):
+        return None
+      body = ", ".join(
+          f"{kt if kt is not None else _plain_repr(k)}:"
+          f" {vt if vt is not None else _plain_repr(item)}"
+          for k, kt, item, vt in walked
+      )
+      return "{" + body + "}"
+    if isinstance(v, (list, tuple, set, frozenset)):
+      elements = [(item, walk(item, depth)) for item in v]
+      if all(text is None for _, text in elements):
+        return None
+      parts = [
+          text if text is not None else _plain_repr(item)
+          for item, text in elements
+      ]
+      if isinstance(v, list):
+        return f"[{', '.join(parts)}]"
+      if isinstance(v, tuple):
+        return f"({parts[0]},)" if len(parts) == 1 else f"({', '.join(parts)})"
+      body = "{" + ", ".join(parts) + "}"
+      return body if isinstance(v, set) else f"frozenset({body})"
+    if (
+        speaks_for_itself
+        and dataclasses.is_dataclass(v)
+        and not isinstance(v, type)
+    ):
+      parts = members(
+          ((f.name, getattr(v, f.name, None)) for f in dataclasses.fields(v)),
+          depth,
+      )
+      return f"{name}({', '.join(parts)})" if parts else None
+    if speaks_for_itself and isinstance(
+        getattr(cls, "model_fields", None), dict
+    ):
+      declared = getattr(v, "__dict__", None) or {}
+      extra = getattr(v, "__pydantic_extra__", None) or {}
+      parts = members(list(declared.items()) + list(extra.items()), depth)
+      return f"{name}({', '.join(parts)})" if parts else None
+    return summarize_object(v, depth)
+
+  def summarize_object(v: Any, depth: int) -> str | None:
+    """Public-attribute summary; private state is inspected but never shown."""
+    held: list[tuple[str, Any]] = []
+    instance_dict = getattr(v, "__dict__", None)
+    if isinstance(instance_dict, dict):
+      held.extend(instance_dict.items())
+    for slot in sorted(public_slot_names(type(v))):
+      try:
+        held.append((slot, getattr(v, slot)))
+      except AttributeError:
+        continue
+    walked = [(name, item, member(name, item, depth)) for name, item in held]
+    if all(text is None for _, _, text in walked):
+      return None
+    parts = [
+        f"{name}={text if text is not None else _plain_repr(item)}"
+        for name, item, text in walked
+        if not name.startswith("_")
+    ]
+    cls_name = type(v).__name__
+    if not parts:
+      return f"<{cls_name}>"
+    return f"<{cls_name} fields={{{', '.join(parts)}}}>"
+
+  return walk(value, 0)
+
+
 def safe_repr(value: Any, caps: Caps) -> str:
-  """``repr(value)`` capped, resilient, with default-form objects summarized."""
+  """``repr(value)`` capped, resilient, credential-masked, defaults summarized."""
   max_len = caps.max_repr_len
   # Fast path: scalars never hit the default-repr regex or summary.
   if type(value) in _SCALAR_TYPES:
@@ -84,17 +306,33 @@ def safe_repr(value: Any, caps: Caps) -> str:
         if len(r) <= max_len
         else r[:max_len] + f"...[{len(r) - max_len} more chars]"
     )
+  if _is_credential_type(type(value)):
+    return f"<{type(value).__name__}>"
   try:
-    r = repr(value)
+    redacted = _redacted_repr(value)
   except Exception as exc:  # pylint: disable=broad-exception-caught
+    # Elided rather than repr'd: the walk stopped partway, so nothing here
+    # says the value is free of secrets.
     logger.warning(
-        "AutoTracingPlugin: repr() failed for %s: %s",
+        "AutoTracingPlugin: redaction failed for %s: %s",
         type(value).__name__,
         exc,
     )
-    r = f"<unrepr-able {type(value).__name__}: {exc!r}>"
-  if _DEFAULT_REPR_RE.match(r):
-    r = _summarize_default(value)
+    return f"<{type(value).__name__} ...>"
+  if redacted is not None:
+    r = redacted
+  else:
+    try:
+      r = repr(value)
+    except Exception as exc:  # pylint: disable=broad-exception-caught
+      logger.warning(
+          "AutoTracingPlugin: repr() failed for %s: %s",
+          type(value).__name__,
+          exc,
+      )
+      r = f"<unrepr-able {type(value).__name__}: {exc!r}>"
+    if _DEFAULT_REPR_RE.match(r):
+      r = _summarize_default(value)
   if len(r) > max_len:
     r = r[:max_len] + f"...[{len(r) - max_len} more chars]"
   return r
@@ -137,6 +375,9 @@ def _summarize_default(value: Any) -> str:
     return f"<{cls}>"
   fields = []
   for k, v in public:
+    if _is_credential_arg_name(k) or _is_credential_type(type(v)):
+      fields.append(f"{k}=<{type(v).__name__}>")
+      continue
     try:
       vr = repr(v)
     except Exception as exc:  # pylint: disable=broad-exception-caught
@@ -174,14 +415,23 @@ def name_value_pairs(
     kwargs: dict[str, Any],
     caps: Caps,
 ) -> list[NamedArg]:
-  """Returns ``[(name, repr)]`` for args + kwargs (no self/cls)."""
+  """Returns ``[(name, repr)]`` for args + kwargs (no self/cls).
+
+  An argument whose name marks it as secret material is dropped outright
+  rather than masked: at the top level the key alone already says the call
+  took a token, and no rendering of the value is worth recording. Values
+  *nested* inside a recorded argument are masked in place instead, because
+  dropping them would misreport the shape of the value that is being traced.
+  """
   pairs: list[NamedArg] = []
   for i, v in enumerate(args):
     name = param_names[i] if i < len(param_names) else f"arg{i}"
-    if name in _SELF_OR_CLS:
+    if name in _SELF_OR_CLS or _is_credential_arg_name(name):
       continue
     pairs.append((name, safe_repr(v, caps)))
   for k, v in kwargs.items():
+    if _is_credential_arg_name(k):
+      continue
     pairs.append((k, safe_repr(v, caps)))
   return pairs
 
@@ -196,6 +446,10 @@ def record_io_on_span(
   """Writes ``adk.fn.*`` attributes onto ``span`` for the call's IO."""
   s = span.set_attribute
   for k, v in pairs:
+    # Repeats the filter in name_value_pairs on purpose: both functions are
+    # public, so pairs may come from a caller that never ran that filter.
+    if _is_credential_arg_name(k):
+      continue
     s(f"adk.fn.arg.{k}", v)
   if exc is not None:
     s("adk.fn.exc_type", type(exc).__qualname__)

src/google/adk/sessions/database_session_service.py

@@ -682,7 +682,11 @@ async def get_session(
         )
 
         if config and config.after_timestamp:
-          after_dt = datetime.fromtimestamp(config.after_timestamp)
+          after_dt = datetime.fromtimestamp(
+              config.after_timestamp, tz=timezone.utc
+          )
+          if self._uses_naive_datetime():
+            after_dt = after_dt.replace(tzinfo=None)
           stmt = stmt.filter(schema.StorageEvent.timestamp >= after_dt)
 
         # Break timestamp ties on id, matching the ordering the stale-session
@@ -932,12 +936,9 @@ async def append_event(self, session: Session, event: Event) -> Event:
           storage_session.state.update(state_deltas["session"])
 
         is_postgresql = self.db_engine.dialect.name == _POSTGRESQL_DIALECT
-        if is_sqlite or is_postgresql:
-          update_time = datetime.fromtimestamp(
-              event.timestamp, timezone.utc
-          ).replace(tzinfo=None)
-        else:
-          update_time = datetime.fromtimestamp(event.timestamp)
+        update_time = datetime.fromtimestamp(event.timestamp, timezone.utc)
+        if self._uses_naive_datetime():
+          update_time = update_time.replace(tzinfo=None)
         storage_session.update_time = update_time
         sql_session.add(schema.StorageEvent.from_event(session, event))
 

src/google/adk/sessions/schemas/shared.py

@@ -31,6 +31,20 @@
 DEFAULT_MAX_VARCHAR_LENGTH = 256
 
 
+def timestamp_to_utc_datetime(timestamp: float) -> datetime.datetime:
+  """Converts a POSIX timestamp to a naive UTC database value."""
+  return datetime.datetime.fromtimestamp(
+      timestamp, tz=datetime.timezone.utc
+  ).replace(tzinfo=None)
+
+
+def utc_datetime_to_timestamp(value: datetime.datetime) -> float:
+  """Converts a database UTC datetime to a POSIX timestamp."""
+  if value.tzinfo is None:
+    value = value.replace(tzinfo=datetime.timezone.utc)
+  return value.timestamp()
+
+
 class DynamicJSON(TypeDecorator[dict[str, Any]]):  # type: ignore[misc]
   """A JSON-like type that uses JSONB on PostgreSQL and TEXT with JSON serialization for other databases."""