Skip to content

Commit ce8d2a3

Browse files
caohy1988claude
andcommitted
feat(sessions): add secret: session state scope (Phase 1)
Introduce a new `secret:` prefix for session state keys that keeps sensitive data (tokens, credentials) in process memory only — never persisted to any storage backend and never logged by BQ Agent Analytics. - Add `State.SECRET_PREFIX` constant and wire it through `extract_state_delta()` so secret keys are excluded from all persistence buckets. - Add process-local cache and lifecycle helpers on `BaseSessionService` (_apply, _trim, _seed, _restore, _evict). - Update all four session services (InMemory, Database, Sqlite, VertexAI) to seed/restore/evict secret state on create/get/delete. - Harden BQ Agent Analytics redaction: redact `secret:*` keys and detect JSON-encoded blobs containing sensitive credential keys. - Accept `secret:` as a valid prefix in instruction template injection. - 32 new tests (unit + integration across all service types). Closes #5112 (Phase 1) Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
1 parent e1913a6 commit ce8d2a3

File tree

10 files changed

+621
-12
lines changed

10 files changed

+621
-12
lines changed

src/google/adk/plugins/bigquery_agent_analytics_plugin.py

Lines changed: 28 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -218,6 +218,24 @@ def _get_tool_origin(tool: "BaseTool") -> str:
218218
})
219219

220220

221+
def _is_sensitive_json_string(value: str) -> bool:
222+
"""Checks if a string is a JSON blob containing sensitive keys.
223+
224+
This catches opaque JSON-encoded strings (e.g. serialized credential
225+
caches) whose inner keys would not be caught by the dict-level
226+
_SENSITIVE_KEYS check.
227+
"""
228+
if not value or value[0] not in ("{", "["):
229+
return False
230+
try:
231+
parsed = json.loads(value)
232+
except (json.JSONDecodeError, ValueError):
233+
return False
234+
if isinstance(parsed, dict):
235+
return bool(_SENSITIVE_KEYS & {k.lower() for k in parsed})
236+
return False
237+
238+
221239
def _recursive_smart_truncate(
222240
obj: Any, max_len: int, seen: Optional[set[int]] = None
223241
) -> tuple[Any, bool]:
@@ -266,10 +284,19 @@ def _recursive_smart_truncate(
266284
for k, v in obj.items():
267285
if isinstance(k, str):
268286
k_lower = k.lower()
269-
if k_lower in _SENSITIVE_KEYS or k_lower.startswith("temp:"):
287+
if (
288+
k_lower in _SENSITIVE_KEYS
289+
or k_lower.startswith("temp:")
290+
or k_lower.startswith("secret:")
291+
):
270292
new_dict[k] = "[REDACTED]"
271293
continue
272294

295+
# Detect JSON-encoded strings that contain sensitive keys.
296+
if isinstance(v, str) and _is_sensitive_json_string(v):
297+
new_dict[k] = "[REDACTED]"
298+
continue
299+
273300
val, trunc = _recursive_smart_truncate(v, max_len, seen)
274301
if trunc:
275302
truncated_any = True

src/google/adk/sessions/_session_util.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,8 @@ def extract_state_delta(
4545
deltas["app"][key.removeprefix(State.APP_PREFIX)] = state[key]
4646
elif key.startswith(State.USER_PREFIX):
4747
deltas["user"][key.removeprefix(State.USER_PREFIX)] = state[key]
48-
elif not key.startswith(State.TEMP_PREFIX):
48+
elif not key.startswith(State.TEMP_PREFIX) and not key.startswith(
49+
State.SECRET_PREFIX
50+
):
4951
deltas["session"][key] = state[key]
5052
return deltas

src/google/adk/sessions/base_session_service.py

Lines changed: 87 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -57,6 +57,22 @@ class BaseSessionService(abc.ABC):
5757
The service provides a set of methods for managing sessions and events.
5858
"""
5959

60+
@property
61+
def _secret_state_cache(
62+
self,
63+
) -> dict[tuple[str, str, str], dict[str, Any]]:
64+
"""Process-local cache for secret-scoped state.
65+
66+
Keyed by (app_name, user_id, session_id).
67+
Lazily initialized to avoid requiring subclasses to call
68+
super().__init__().
69+
"""
70+
try:
71+
return self.__secret_state_cache
72+
except AttributeError:
73+
self.__secret_state_cache: dict[tuple[str, str, str], dict[str, Any]] = {}
74+
return self.__secret_state_cache
75+
6076
@abc.abstractmethod
6177
async def create_session(
6278
self,
@@ -120,6 +136,10 @@ async def append_event(self, session: Session, event: Event) -> Event:
120136
# read temp values (e.g. output_key='temp:my_key' in SequentialAgent).
121137
self._apply_temp_state(session, event)
122138
event = self._trim_temp_delta_state(event)
139+
# Apply secret-scoped state to in-memory session and process cache
140+
# BEFORE trimming, so the session retains secret values across turns.
141+
self._apply_secret_state(session, event)
142+
event = self._trim_secret_delta_state(event)
123143
self._update_session_state(session, event)
124144
session.events.append(event)
125145
return event
@@ -154,6 +174,73 @@ def _trim_temp_delta_state(self, event: Event) -> Event:
154174
}
155175
return event
156176

177+
def _apply_secret_state(self, session: Session, event: Event) -> None:
178+
"""Applies secret-scoped state to in-memory session and process cache.
179+
180+
Secret state survives across turns (via the process-local cache) but
181+
is never persisted to storage. The event delta is trimmed separately
182+
by _trim_secret_delta_state.
183+
"""
184+
if not event.actions or not event.actions.state_delta:
185+
return
186+
cache_key = (session.app_name, session.user_id, session.id)
187+
for key, value in event.actions.state_delta.items():
188+
if key.startswith(State.SECRET_PREFIX):
189+
session.state[key] = value
190+
self._secret_state_cache.setdefault(cache_key, {})[key] = value
191+
192+
def _trim_secret_delta_state(self, event: Event) -> Event:
193+
"""Removes secret-scoped keys from event delta before persistence."""
194+
if not event.actions or not event.actions.state_delta:
195+
return event
196+
event.actions.state_delta = {
197+
key: value
198+
for key, value in event.actions.state_delta.items()
199+
if not key.startswith(State.SECRET_PREFIX)
200+
}
201+
return event
202+
203+
def _seed_secret_state_on_create(
204+
self,
205+
*,
206+
app_name: str,
207+
user_id: str,
208+
session_id: str,
209+
state: Optional[dict[str, Any]],
210+
) -> Optional[dict[str, Any]]:
211+
"""Extracts secret-scoped keys from initial state into the cache.
212+
213+
Returns the state dict with secret keys removed (for persistence)
214+
but seeds them in the process-local cache so get_session() can
215+
restore them.
216+
"""
217+
if not state:
218+
return state
219+
secret_keys = {
220+
k: v for k, v in state.items() if k.startswith(State.SECRET_PREFIX)
221+
}
222+
if not secret_keys:
223+
return state
224+
cache_key = (app_name, user_id, session_id)
225+
self._secret_state_cache.setdefault(cache_key, {}).update(secret_keys)
226+
return {
227+
k: v for k, v in state.items() if not k.startswith(State.SECRET_PREFIX)
228+
}
229+
230+
def _restore_secret_state(self, session: Session) -> None:
231+
"""Merges cached secret state into an in-memory session."""
232+
cache_key = (session.app_name, session.user_id, session.id)
233+
secret_state = self._secret_state_cache.get(cache_key, {})
234+
for key, value in secret_state.items():
235+
session.state[key] = value
236+
237+
def _evict_secret_state(
238+
self, app_name: str, user_id: str, session_id: str
239+
) -> None:
240+
"""Removes cached secret state for a deleted session."""
241+
cache_key = (app_name, user_id, session_id)
242+
self._secret_state_cache.pop(cache_key, None)
243+
157244
def _update_session_state(self, session: Session, event: Event) -> None:
158245
"""Updates the session state based on the event."""
159246
if not event.actions or not event.actions.state_delta:

src/google/adk/sessions/database_session_service.py

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -445,6 +445,21 @@ async def create_session(
445445
defaults={"app_name": app_name, "user_id": user_id, "state": {}},
446446
)
447447

448+
# Extract secret keys before they are dropped by
449+
# extract_state_delta; we'll seed the cache after the session
450+
# is committed and has a final ID.
451+
secret_keys = {}
452+
if state:
453+
secret_keys = {
454+
k: v for k, v in state.items() if k.startswith(State.SECRET_PREFIX)
455+
}
456+
if secret_keys:
457+
state = {
458+
k: v
459+
for k, v in state.items()
460+
if not k.startswith(State.SECRET_PREFIX)
461+
}
462+
448463
# Extract state deltas
449464
state_deltas = _session_util.extract_state_delta(state)
450465
app_state_delta = state_deltas["app"]
@@ -482,6 +497,13 @@ async def create_session(
482497
session = storage_session.to_session(
483498
state=merged_state, is_sqlite=is_sqlite
484499
)
500+
501+
# Seed secret state into the cache now that session.id is resolved.
502+
if secret_keys:
503+
cache_key = (app_name, user_id, session.id)
504+
self._secret_state_cache.setdefault(cache_key, {}).update(secret_keys)
505+
for key, value in secret_keys.items():
506+
session.state[key] = value
485507
return session
486508

487509
@override
@@ -547,6 +569,7 @@ async def get_session(
547569
session = storage_session.to_session(
548570
state=merged_state, events=events, is_sqlite=is_sqlite
549571
)
572+
self._restore_secret_state(session)
550573
return session
551574

552575
@override
@@ -615,6 +638,7 @@ async def delete_session(
615638
)
616639
await sql_session.execute(stmt)
617640
await sql_session.commit()
641+
self._evict_secret_state(app_name, user_id, session_id)
618642

619643
@override
620644
async def append_event(self, session: Session, event: Event) -> Event:
@@ -627,6 +651,9 @@ async def append_event(self, session: Session, event: Event) -> Event:
627651
self._apply_temp_state(session, event)
628652
# Trim temp state before persisting
629653
event = self._trim_temp_delta_state(event)
654+
# Apply secret state to in-memory session and process cache.
655+
self._apply_secret_state(session, event)
656+
event = self._trim_secret_delta_state(event)
630657

631658
# 1. Validate the session has not gone stale.
632659
# 2. Update session attributes based on event config.

src/google/adk/sessions/in_memory_session_service.py

Lines changed: 23 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -118,6 +118,22 @@ def _create_session_impl(
118118
app_name=app_name, user_id=user_id, session_id=session_id
119119
):
120120
raise AlreadyExistsError(f'Session with id {session_id} already exists.')
121+
122+
session_id = (
123+
session_id.strip()
124+
if session_id and session_id.strip()
125+
else platform_uuid.new_uuid()
126+
)
127+
128+
# Seed secret state into the process-local cache before
129+
# extract_state_delta, which would otherwise drop secret keys.
130+
state = self._seed_secret_state_on_create(
131+
app_name=app_name,
132+
user_id=user_id,
133+
session_id=session_id,
134+
state=state,
135+
)
136+
121137
state_deltas = _session_util.extract_state_delta(state)
122138
app_state_delta = state_deltas['app']
123139
user_state_delta = state_deltas['user']
@@ -129,11 +145,6 @@ def _create_session_impl(
129145
user_state_delta
130146
)
131147

132-
session_id = (
133-
session_id.strip()
134-
if session_id and session_id.strip()
135-
else platform_uuid.new_uuid()
136-
)
137148
session = Session(
138149
app_name=app_name,
139150
user_id=user_id,
@@ -149,7 +160,9 @@ def _create_session_impl(
149160
self.sessions[app_name][user_id][session_id] = session
150161

151162
copied_session = _copy_session(session)
152-
return self._merge_state(app_name, user_id, copied_session)
163+
merged = self._merge_state(app_name, user_id, copied_session)
164+
self._restore_secret_state(merged)
165+
return merged
153166

154167
@override
155168
async def get_session(
@@ -219,7 +232,9 @@ def _get_session_impl(
219232
copied_session.events = copied_session.events[i + 1 :]
220233

221234
# Return a copy of the session object with merged state.
222-
return self._merge_state(app_name, user_id, copied_session)
235+
merged = self._merge_state(app_name, user_id, copied_session)
236+
self._restore_secret_state(merged)
237+
return merged
223238

224239
def _merge_state(
225240
self, app_name: str, user_id: str, copied_session: Session
@@ -311,6 +326,7 @@ def _delete_session_impl(
311326
return
312327

313328
self.sessions[app_name][user_id].pop(session_id)
329+
self._evict_secret_state(app_name, user_id, session_id)
314330

315331
@override
316332
async def append_event(self, session: Session, event: Event) -> Event:

src/google/adk/sessions/sqlite_session_service.py

Lines changed: 19 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -179,6 +179,15 @@ async def create_session(
179179
f"Session with id {session_id} already exists."
180180
)
181181

182+
# Seed secret state into the process-local cache before
183+
# extract_state_delta, which would otherwise drop secret keys.
184+
state = self._seed_secret_state_on_create(
185+
app_name=app_name,
186+
user_id=user_id,
187+
session_id=session_id,
188+
state=state,
189+
)
190+
182191
# Extract state deltas
183192
state_deltas = _session_util.extract_state_delta(state)
184193
app_state_delta = state_deltas["app"]
@@ -218,14 +227,16 @@ async def create_session(
218227
merged_state = _merge_state(
219228
storage_app_state, storage_user_state, session_state
220229
)
221-
return Session(
230+
session = Session(
222231
app_name=app_name,
223232
user_id=user_id,
224233
id=session_id,
225234
state=merged_state,
226235
events=[],
227236
last_update_time=now,
228237
)
238+
self._restore_secret_state(session)
239+
return session
229240

230241
@override
231242
async def get_session(
@@ -284,14 +295,16 @@ async def get_session(
284295
for event_data in reversed(storage_events_data)
285296
]
286297

287-
return Session(
298+
session = Session(
288299
app_name=app_name,
289300
user_id=user_id,
290301
id=session_id,
291302
state=merged_state,
292303
events=events,
293304
last_update_time=last_update_time,
294305
)
306+
self._restore_secret_state(session)
307+
return session
295308

296309
@override
297310
async def list_sessions(
@@ -358,6 +371,7 @@ async def delete_session(
358371
(app_name, user_id, session_id),
359372
)
360373
await db.commit()
374+
self._evict_secret_state(app_name, user_id, session_id)
361375

362376
@override
363377
async def append_event(self, session: Session, event: Event) -> Event:
@@ -369,6 +383,9 @@ async def append_event(self, session: Session, event: Event) -> Event:
369383
self._apply_temp_state(session, event)
370384
# Trim temp state before persisting
371385
event = self._trim_temp_delta_state(event)
386+
# Apply secret state to in-memory session and process cache.
387+
self._apply_secret_state(session, event)
388+
event = self._trim_secret_delta_state(event)
372389
event_timestamp = event.timestamp
373390

374391
async with self._get_db_connection() as db:

src/google/adk/sessions/state.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@ class State:
2323
APP_PREFIX = "app:"
2424
USER_PREFIX = "user:"
2525
TEMP_PREFIX = "temp:"
26+
SECRET_PREFIX = "secret:"
2627

2728
def __init__(self, value: dict[str, Any], delta: dict[str, Any]):
2829
"""

0 commit comments

Comments
 (0)