feat: add nested state access support with dot and bracket notation

- Add support for accessing nested dictionary structures in state variables
- Support dot notation (state.key.subkey) and bracket notation (state['key']['subkey'])
- Support mixed notation (state.key['subkey'])
- Maintain backward compatibility with existing flat state access
- Add comprehensive tests for all nested access patterns
- Handle optional syntax for nested paths
- Proper error handling for missing nested keys

Functions added:
- _parse_nested_path(): Parse variable paths with dot/bracket notation
- _get_nested_value(): Safely traverse nested dictionary structures
- _is_valid_state_name_or_nested(): Extended validation for nested paths

All existing tests pass, ensuring backward compatibility is maintained.

Author: FingerLiu

src/google/adk/utils/instructions_utils.py

@@ -93,19 +93,97 @@ async def _replace_match(match) -> str:
         raise KeyError(f'Artifact {var_name} not found.')
       return str(artifact)
     else:
-      if not _is_valid_state_name(var_name):
+      if not _is_valid_state_name_or_nested(var_name):
         return match.group()
-      if var_name in invocation_context.session.state:
-        return str(invocation_context.session.state[var_name])
-      else:
+      
+      try:
+        keys = _parse_nested_path(var_name)
+        value = _get_nested_value(invocation_context.session.state, keys)
+        return str(value)
+      except KeyError as e:
         if optional:
           return ''
         else:
-          raise KeyError(f'Context variable not found: `{var_name}`.')
+          raise KeyError(f'Context variable not found: `{var_name}`. {str(e)}')
 
   return await _async_sub(r'{+[^{}]*}+', _replace_match, template)
 
 
+def _parse_nested_path(var_name: str) -> list[str]:
+  """Parse a nested variable path into individual keys.
+  
+  Supports both dot notation (key.subkey) and bracket notation (key['subkey']).
+  Mixed notation is also supported (key.subkey['nested']).
+  
+  Args:
+    var_name: The variable name to parse (e.g., "user.profile.name" or "user['profile']['name']")
+    
+  Returns:
+    List of keys to traverse the nested structure.
+  """
+  if '.' not in var_name and '[' not in var_name:
+    return [var_name]
+  
+  keys = []
+  current_key = ""
+  i = 0
+  
+  while i < len(var_name):
+    char = var_name[i]
+    
+    if char == '.':
+      if current_key:
+        keys.append(current_key)
+        current_key = ""
+    elif char == '[':
+      if current_key:
+        keys.append(current_key)
+        current_key = ""
+      bracket_end = var_name.find(']', i)
+      if bracket_end == -1:
+        raise ValueError(f"Unclosed bracket in variable name: {var_name}")
+      
+      bracket_content = var_name[i+1:bracket_end]
+      if (bracket_content.startswith('"') and bracket_content.endswith('"')) or \
+         (bracket_content.startswith("'") and bracket_content.endswith("'")):
+        bracket_content = bracket_content[1:-1]
+      
+      keys.append(bracket_content)
+      i = bracket_end
+    else:
+      current_key += char
+    
+    i += 1
+  
+  if current_key:
+    keys.append(current_key)
+  
+  return keys
+
+
+def _get_nested_value(data: dict, keys: list[str]):
+  """Get a value from nested dictionary structure using a list of keys.
+  
+  Args:
+    data: The dictionary to traverse
+    keys: List of keys to traverse the nested structure
+    
+  Returns:
+    The value at the nested path
+    
+  Raises:
+    KeyError: If any key in the path doesn't exist
+  """
+  current = data
+  for key in keys:
+    if not isinstance(current, dict):
+      raise KeyError(f"Cannot access key '{key}' on non-dict value")
+    if key not in current:
+      raise KeyError(f"Key '{key}' not found")
+    current = current[key]
+  return current
+
+
 def _is_valid_state_name(var_name):
   """Checks if the variable name is a valid state name.
 
@@ -129,3 +207,32 @@ def _is_valid_state_name(var_name):
     if (parts[0] + ':') in prefixes:
       return parts[1].isidentifier()
   return False
+
+
+def _is_valid_state_name_or_nested(var_name: str) -> bool:
+  """Checks if the variable name is a valid state name or nested path.
+  
+  Valid state is either:
+    - Valid identifier (existing behavior)
+    - <Valid prefix>:<Valid identifier> (existing behavior)  
+    - Nested path with dot notation (key.subkey.nested)
+    - Nested path with bracket notation (key['subkey']['nested'])
+    - Mixed notation (key.subkey['nested'])
+  
+  Args:
+    var_name: The variable name to check.
+    
+  Returns:
+    True if the variable name is valid, False otherwise.
+  """
+  if _is_valid_state_name(var_name):
+    return True
+  
+  try:
+    keys = _parse_nested_path(var_name)
+    for key in keys:
+      if not _is_valid_state_name(key):
+        return False
+    return len(keys) > 1
+  except ValueError:
+    return False

tests/unittests/utils/test_instructions_utils.py

@@ -214,3 +214,113 @@ async def test_inject_session_state_artifact_service_not_initialized_raises_valu
     await instructions_utils.inject_session_state(
         instruction_template, invocation_context
     )
+
+
+@pytest.mark.asyncio
+async def test_inject_session_state_with_nested_dot_notation():
+  instruction_template = "User name: {user.profile.name}, Age: {user.profile.age}"
+  invocation_context = await _create_test_readonly_context(
+      state={
+          "user": {
+              "profile": {
+                  "name": "John Doe",
+                  "age": 30
+              }
+          }
+      }
+  )
+
+  populated_instruction = await instructions_utils.inject_session_state(
+      instruction_template, invocation_context
+  )
+  assert populated_instruction == "User name: John Doe, Age: 30"
+
+
+@pytest.mark.asyncio
+async def test_inject_session_state_with_nested_bracket_notation():
+  instruction_template = "User name: {user['profile']['name']}, City: {user['address']['city']}"
+  invocation_context = await _create_test_readonly_context(
+      state={
+          "user": {
+              "profile": {"name": "Jane Smith"},
+              "address": {"city": "New York"}
+          }
+      }
+  )
+
+  populated_instruction = await instructions_utils.inject_session_state(
+      instruction_template, invocation_context
+  )
+  assert populated_instruction == "User name: Jane Smith, City: New York"
+
+
+@pytest.mark.asyncio
+async def test_inject_session_state_with_mixed_notation():
+  instruction_template = "Data: {config.database['host']}, Port: {config['database'].port}"
+  invocation_context = await _create_test_readonly_context(
+      state={
+          "config": {
+              "database": {
+                  "host": "localhost",
+                  "port": 5432
+              }
+          }
+      }
+  )
+
+  populated_instruction = await instructions_utils.inject_session_state(
+      instruction_template, invocation_context
+  )
+  assert populated_instruction == "Data: localhost, Port: 5432"
+
+
+@pytest.mark.asyncio
+async def test_inject_session_state_with_nested_optional():
+  instruction_template = "Optional nested: {user.profile.nickname?}"
+  invocation_context = await _create_test_readonly_context(
+      state={"user": {"profile": {"name": "John"}}}
+  )
+
+  populated_instruction = await instructions_utils.inject_session_state(
+      instruction_template, invocation_context
+  )
+  assert populated_instruction == "Optional nested: "
+
+
+@pytest.mark.asyncio
+async def test_inject_session_state_with_nested_missing_raises_error():
+  instruction_template = "Missing: {user.profile.missing}"
+  invocation_context = await _create_test_readonly_context(
+      state={"user": {"profile": {"name": "John"}}}
+  )
+
+  with pytest.raises(KeyError, match="Context variable not found: `user.profile.missing`"):
+    await instructions_utils.inject_session_state(
+        instruction_template, invocation_context
+    )
+
+
+@pytest.mark.asyncio
+async def test_inject_session_state_with_invalid_nested_path():
+  instruction_template = "Invalid: {user.profile[unclosed}"
+  invocation_context = await _create_test_readonly_context(
+      state={"user": {"profile": {"name": "John"}}}
+  )
+
+  populated_instruction = await instructions_utils.inject_session_state(
+      instruction_template, invocation_context
+  )
+  assert populated_instruction == "Invalid: {user.profile[unclosed}"
+
+
+@pytest.mark.asyncio
+async def test_inject_session_state_backward_compatibility():
+  instruction_template = "Hello {user_name}, status: {app_status}"
+  invocation_context = await _create_test_readonly_context(
+      state={"user_name": "Alice", "app_status": "active"}
+  )
+
+  populated_instruction = await instructions_utils.inject_session_state(
+      instruction_template, invocation_context
+  )
+  assert populated_instruction == "Hello Alice, status: active"