fix: Fix instructions_utils matching invalid nested paths

Co-authored-by: Bo Yang <ybo@google.com>
PiperOrigin-RevId: 934582581

Author: boyangsvl

src/google/adk/utils/instructions_utils.py

@@ -120,11 +120,28 @@ def _get_nested_value(obj: Any, path: str) -> Any:
 
     return current
 
-  async def _replace_match(match) -> str:
-    full_path = match.group().lstrip('{').rstrip('}').strip()
+  def _is_valid_path(path: str) -> bool:
+    """Checks if the path is a valid state variable path."""
+    parts = path.split('.')
+    if not parts:
+      return False
+
+    # Check first segment (can have prefix)
+    first_seg = parts[0].removesuffix('?')
+    if not _is_valid_state_name(first_seg):
+      return False
+
+    # Check subsequent segments (must be plain identifiers)
+    for part in parts[1:]:
+      seg = part.removesuffix('?')
+      if not seg.isidentifier():
+        return False
+
+    return True
 
-    if full_path.startswith('artifact.'):
-      var_name = full_path.removeprefix('artifact.')
+  async def _evaluate_path(path: str) -> str:
+    if path.startswith('artifact.'):
+      var_name = path.removeprefix('artifact.')
       optional = var_name.endswith('?')
       if optional:
         var_name = var_name.removesuffix('?')
@@ -147,17 +164,40 @@ async def _replace_match(match) -> str:
           raise KeyError(f'Artifact {var_name} not found.')
       return str(artifact)
     else:
-      if not _is_valid_state_name(full_path.split('.')[0].removesuffix('?')):
-        return match.group()
-
       try:
-        value = _get_nested_value(invocation_context.session.state, full_path)
+        value = _get_nested_value(invocation_context.session.state, path)
 
         if value is None:
           return ''
         return str(value)
       except KeyError as e:
-        raise KeyError(f'Context variable not found: `{full_path}`.') from e
+        raise KeyError(f'Context variable not found: `{path}`.') from e
+
+  async def _replace_match(match) -> str:
+    raw_match = match.group()
+    leading_count = len(raw_match) - len(raw_match.lstrip('{'))
+    trailing_count = len(raw_match) - len(raw_match.rstrip('}'))
+    full_path = raw_match.lstrip('{').rstrip('}').strip()
+
+    if leading_count == trailing_count:
+      n = leading_count
+      if n % 2 == 0:
+        # Even: Escaped, no evaluation
+        half_n = n // 2
+        return '{' * half_n + full_path + '}' * half_n
+      else:
+        # Odd: Evaluate and wrap
+        if not _is_valid_path(full_path) and not full_path.startswith('artifact.'):
+          return raw_match
+        evaluated_value = await _evaluate_path(full_path)
+        wrap_braces = (n - 1) // 2
+        return '{' * wrap_braces + evaluated_value + '}' * wrap_braces
+    else:
+      # Asymmetric: fallback to old behavior (treat as N=1 if valid path)
+      if not _is_valid_path(full_path) and not full_path.startswith('artifact.'):
+        return raw_match
+      return await _evaluate_path(full_path)
+
 
   return await _async_sub(r'{+[^{}]*}+', _replace_match, template)
 

tests/unittests/utils/test_instructions_utils.py

@@ -497,3 +497,88 @@ def __init__(self):
       instruction_template, invocation_context
   )
   assert populated_instruction == "Name: Frank, Role: Engineer"
+
+
+@pytest.mark.asyncio
+async def test_inject_session_state_with_invalid_nested_path_returns_original():
+  """Test that invalid nested paths (e.g. containing spaces or equals) are ignored."""
+  instruction_template = "Value: {LabelConfidence.Enum confidence = 441216274;}"
+  invocation_context = await _create_test_readonly_context()
+
+  populated_instruction = await instructions_utils.inject_session_state(
+      instruction_template, invocation_context
+  )
+  assert (
+      populated_instruction
+      == "Value: {LabelConfidence.Enum confidence = 441216274;}"
+  )
+
+
+@pytest.mark.asyncio
+async def test_inject_session_state_escaped_braces():
+  instruction_template = "This is a literal {{placeholder}} and this is {value}."
+  invocation_context = await _create_test_readonly_context(
+      state={"value": "real_value"}
+  )
+  populated_instruction = await instructions_utils.inject_session_state(
+      instruction_template, invocation_context
+  )
+  assert populated_instruction == "This is a literal {placeholder} and this is real_value."
+
+
+@pytest.mark.asyncio
+async def test_inject_session_state_escaped_braces_missing_state():
+  instruction_template = "This is a literal {{task_N.result}}."
+  # task_N is NOT in state
+  invocation_context = await _create_test_readonly_context(state={})
+  populated_instruction = await instructions_utils.inject_session_state(
+      instruction_template, invocation_context
+  )
+  assert populated_instruction == "This is a literal {task_N.result}."
+
+
+@pytest.mark.asyncio
+async def test_inject_session_state_triple_braces_evaluates_and_wraps():
+  instruction_template = "This is a wrapped value: {{{value}}}."
+  invocation_context = await _create_test_readonly_context(
+      state={"value": "real_value"}
+  )
+  populated_instruction = await instructions_utils.inject_session_state(
+      instruction_template, invocation_context
+  )
+  assert populated_instruction == "This is a wrapped value: {real_value}."
+
+
+@pytest.mark.asyncio
+async def test_inject_session_state_quadruple_braces():
+  instruction_template = "This is literal double braces: {{{{placeholder}}}}."
+  invocation_context = await _create_test_readonly_context(state={})
+  populated_instruction = await instructions_utils.inject_session_state(
+      instruction_template, invocation_context
+  )
+  assert populated_instruction == "This is literal double braces: {{placeholder}}."
+
+
+@pytest.mark.asyncio
+async def test_inject_session_state_asymmetric_braces_left():
+  instruction_template = "Asymmetric left: {{value}."
+  invocation_context = await _create_test_readonly_context(
+      state={"value": "real_value"}
+  )
+  populated_instruction = await instructions_utils.inject_session_state(
+      instruction_template, invocation_context
+  )
+  assert populated_instruction == "Asymmetric left: real_value."
+
+
+@pytest.mark.asyncio
+async def test_inject_session_state_asymmetric_braces_right():
+  instruction_template = "Asymmetric right: {value}}."
+  invocation_context = await _create_test_readonly_context(
+      state={"value": "real_value"}
+  )
+  populated_instruction = await instructions_utils.inject_session_state(
+      instruction_template, invocation_context
+  )
+  assert populated_instruction == "Asymmetric right: real_value."
+