Lab 03 — Sessions, Scoped State & the Callback Chain

Phase 22 · Lab 03 · Phase README · Warmup

The problem

Two Google ADK subsystems make an agent stateful and controllable. This lab builds faithful miniatures of both.

1) Sessions & scoped state. A SessionService (InMemory / Database / VertexAI in real ADK) owns Session objects, and every session has a state dict — the agent's working memory across turns. The trick is state scope prefixes: a key's prefix decides how widely it is shared:

session.state is a single dict-like view that routes each read/write to the right backing store by prefix. This is Phase 04 context-engineering/memory, productized: ADK separates "what this chat knows" from "what this user always wants" from "what the whole app shares" without you managing three dictionaries by hand.

2) The callback chain. Around each step of an agent, ADK fires paired callbacks you register: before_agent/after_agent, before_model/after_model, before_tool/after_tool. They are the interception layer — logging, metrics, caching, and guardrails (Phase 10). The load-bearing rule: a before_* callback that returns a value SHORT-CIRCUITS the wrapped step. Return a response from before_model and the (paid) model call is skipped — a cache or a policy block. Return a result from before_tool and the tool never runs — a guardrail denying a dangerous action. An after_* callback can transform the step's output (redact PII, clamp a value). Returning None means "proceed / don't modify".

Everything is injected and deterministic: the model is a pure policy, session ids are a counter (no uuid), no clock. That is precisely how you unit-test ADK guardrails — assert a blocked tool never executed and that a cache hit skipped the model.

What you build

File map

Run it

pytest test_lab.py -v                       # your lab.py (red until you implement)
LAB_MODULE=solution pytest test_lab.py -v   # the reference (green)
python solution.py                          # worked example

Success criteria

• Unprefixed state is private to a session; user: keys are shared across that user's sessions (but not other users); app: keys are shared across all users of the app (but not other apps). Seeding initial state routes by prefix too.
• The six callbacks fire in order: before_agent → before_model → after_model → after_agent (with before_tool/after_tool wrapping any tool call).
• A before_model_callback returning a response skips the model (cache hit); a before_tool_callback returning a value blocks the tool (the dangerous function never runs); a before_agent_callback returning a value skips the whole agent body.
• after_* callbacks transform output (uppercase a response, redact a tool result, wrap a final answer). Returning None proceeds unchanged.
• All 25 tests pass under both lab and solution.

How this maps to the real stack

• State, Session, and InMemorySessionService are real ADK classes in google.adk.sessions. The scope prefixes (user:, app:, temp:) are the real ADK convention, and session.state really is a view that routes writes to the right tier; in production you swap InMemorySessionService for DatabaseSessionService or VertexAiSessionService and the same scoping semantics persist to a real store. Long-term recall across sessions is a separate MemoryService (cross-reference Phase 04).
• The six callbacks are the real ADK hook names (before_agent_callback, after_agent_callback, before_model_callback, after_model_callback, before_tool_callback, after_tool_callback), registered on an agent, and the short-circuit contract is exactly real: return a Content from before_model_callback and ADK skips the model; return a dict from before_tool_callback and ADK skips the tool. This is the canonical place to implement guardrails (Phase 10), response caching (Phase 14), and PII redaction. Real callbacks receive a CallbackContext/ToolContext with .state and more; ours carries .state and agent_name.
• Writing the final response to session.state[output_key] is the same behavior you saw in Lab 01 and relied on in Lab 02 — the through-line of the phase.

Limits of the miniature. No persistence (in-memory only), temp: is per-session rather than strictly per-invocation, no async, and the model/tool loop is simplified (one tool round per model turn, capped by max_steps). The two mechanisms that matter — prefix-scoped state sharing and the before_* short-circuit — are faithful.

Extensions (your own machine)

• Add a DatabaseSessionService backed by sqlite3: same State view, persisted stores, and prove user:/app: state survives a process restart.
• Implement a real response cache in before_model_callback keyed by a hash of the request, and measure the model-call reduction on a repeated-query workload (Phase 14).
• Build an injection guardrail as a before_tool_callback that scans arguments for an exfil URL or a path outside an allow-list and blocks the call (Phase 10), with a test that the tool never executed.
• Add a MemoryService stub and an after_agent_callback that writes the turn's summary to long-term memory for later retrieval.

Interview / resume signal

"Built Google ADK's stateful/controllable core: an in-memory SessionService with prefix-scoped state (user:/app:/temp: sharing tiers) and the six-callback lifecycle chain — implementing the before_* short-circuit as a cache (skip the paid model) and a guardrail (block a dangerous tool before it runs), with after_* transforms for PII redaction — all deterministic and unit-tested."