diff --git a/src/google/adk/sessions/state.py b/src/google/adk/sessions/state.py index 587bfa16e1..ee37b5699b 100644 --- a/src/google/adk/sessions/state.py +++ b/src/google/adk/sessions/state.py @@ -85,10 +85,40 @@ def __init__( self._delta = delta self._schema = schema + def _resolve_temp_fallback(self, key: str) -> tuple[bool, Any]: + """Resolves ephemeral temp: value if present for a bare key lookup.""" + if not isinstance(key, str) or key.startswith(self.TEMP_PREFIX): + return False, None + temp_key = f"{self.TEMP_PREFIX}{key}" + if temp_key in self._delta: + return True, self._delta[temp_key] + if temp_key in self._value: + return True, self._value[temp_key] + return False, None + def __getitem__(self, key: str) -> Any: """Returns the value of the state dict for the given key.""" + has_temp, temp_val = self._resolve_temp_fallback(key) if key in self._delta: - return self._delta[key] + val = self._delta[key] + if ( + has_temp + and isinstance(val, str) + and val.startswith("[REDACTED_SECRET:") + ): + return temp_val + return val + if key in self._value: + val = self._value[key] + if ( + has_temp + and isinstance(val, str) + and val.startswith("[REDACTED_SECRET:") + ): + return temp_val + return val + if has_temp: + return temp_val return self._value[key] def __setitem__(self, key: str, value: Any) -> None: @@ -100,9 +130,20 @@ def __setitem__(self, key: str, value: Any) -> None: self._value[key] = value self._delta[key] = value + def set_ephemeral_value(self, key: str, value: Any) -> None: + """Sets an in-memory value without adding it to the persistent delta.""" + if self._schema is not None and isinstance(self._schema, type): + _validate_state_entry(self._schema, key, value) + self._value[key] = value + def __contains__(self, key: object) -> bool: """Whether the state dict contains the given key.""" - return key in self._value or key in self._delta + if key in self._value or key in self._delta: + return True + if isinstance(key, str): + has_temp, _ = self._resolve_temp_fallback(key) + return has_temp + return False def setdefault(self, key: str, default: Any = None) -> Any: """Gets the value of a key, or sets it to a default if the key doesn't exist.""" @@ -135,4 +176,14 @@ def to_dict(self) -> dict[str, Any]: result: dict[str, Any] = {} result.update(self._value) result.update(self._delta) + for k, v in list(result.items()): + if isinstance(k, str) and k.startswith(self.TEMP_PREFIX): + bare = k[len(self.TEMP_PREFIX) :] + if bare: + existing = result.get(bare) + if existing is None or ( + isinstance(existing, str) + and existing.startswith("[REDACTED_SECRET:") + ): + result[bare] = v return result