feat(models): surface and recover environment_id from interactions

Add environment_id to LlmResponse, populated from the interactions stream loop,
and return it alongside interaction_id from the previous-interaction lookup so
stateful agents can reuse a sandbox environment across turns.

Co-authored-by: Haran Rajkumar <haranrk@google.com>
PiperOrigin-RevId: 940034114

Author: haranrk

src/google/adk/flows/llm_flows/interactions_processor.py

@@ -37,16 +37,17 @@ def _is_event_in_branch(current_branch: Optional[str], event: Event) -> bool:
   return event.branch == current_branch or not event.branch
 
 
-def _find_previous_interaction_id(
+def _find_previous_interaction_state(
     events: list[Event],
     *,
     agent_name: str,
     current_branch: Optional[str],
-) -> Optional[str]:
-  """Find the most recent interaction_id authored by ``agent_name``.
+) -> tuple[Optional[str], Optional[str]]:
+  """Find the most recent (interaction_id, environment_id) for ``agent_name``.
 
   Scans ``events`` in reverse, skipping events outside ``current_branch``, and
-  returns the first ``interaction_id`` from an event authored by this agent.
+  returns the ids from the first event authored by this agent that carries an
+  interaction_id.
   """
   logger.debug(
       'Finding previous_interaction_id: agent=%s, branch=%s, num_events=%d',
@@ -75,8 +76,8 @@ def _find_previous_interaction_id(
           agent_name,
           event.interaction_id,
       )
-      return event.interaction_id
-  return None
+      return event.interaction_id, event.environment_id
+  return None, None
 
 
 class InteractionsRequestProcessor(BaseLlmRequestProcessor):
@@ -126,11 +127,12 @@ def _find_previous_interaction_id(
       self, invocation_context: 'InvocationContext'
   ) -> Optional[str]:
     """Find the previous interaction ID from session events."""
-    return _find_previous_interaction_id(
+    interaction_id, _ = _find_previous_interaction_state(
         invocation_context.session.events,
         agent_name=invocation_context.agent.name,
         current_branch=invocation_context.branch,
     )
+    return interaction_id
 
 
 # Module-level processor instance for use in flow configuration

src/google/adk/models/interactions_utils.py

@@ -112,6 +112,29 @@ def _extract_stream_interaction_id(
   return None
 
 
+def _extract_stream_environment_id(
+    event: InteractionSSEEvent,
+) -> str | None:
+  """Extract the environment id from an Interactions SSE event, if present.
+
+  The non-streaming ``Interaction`` declares an ``environment_id`` field. On
+  streaming SSE events the id is read opportunistically from the carried
+  interaction (created/completed events allow extra fields), so it is returned
+  only when the API actually includes it and is ``None`` otherwise.
+  """
+  interaction = None
+  if isinstance(event, (InteractionCreatedEvent, InteractionCompletedEvent)):
+    interaction = event.interaction
+  elif isinstance(event, Interaction):
+    interaction = event
+
+  if interaction is None:
+    return None
+
+  env_id = getattr(interaction, 'environment_id', None)
+  return env_id if isinstance(env_id, str) else None
+
+
 def _encode_base64_string(data: bytes) -> str:
   """Encode bytes to a base64 string."""
   return base64.b64encode(data).decode('utf-8')
@@ -1445,6 +1468,7 @@ async def _create_interactions(
     LlmResponse objects converted from interaction responses.
   """
   current_interaction_id: str | None = None
+  current_environment_id: str | None = None
 
   if stream:
     responses = await api_client.aio.interactions.create(
@@ -1456,18 +1480,24 @@ async def _create_interactions(
       interaction_id = _extract_stream_interaction_id(event)
       if interaction_id:
         current_interaction_id = interaction_id
+      environment_id = _extract_stream_environment_id(event)
+      if environment_id:
+        current_environment_id = environment_id
       llm_response = convert_interaction_event_to_llm_response(
           event, state, current_interaction_id
       )
       if llm_response:
+        llm_response.environment_id = current_environment_id
         yield llm_response
   else:
     interaction = await api_client.aio.interactions.create(
         **create_kwargs, stream=False
     )
     logger.info('Interaction response received.')
     logger.debug(build_interactions_response_log(interaction))
-    yield convert_interaction_to_llm_response(interaction)
+    llm_response = convert_interaction_to_llm_response(interaction)
+    llm_response.environment_id = interaction.environment_id
+    yield llm_response
 
 
 async def generate_content_via_interactions(

src/google/adk/models/llm_response.py

@@ -155,6 +155,14 @@ class LlmResponse(BaseModel):
   It can be used to identify and chain interactions for stateful conversations.
   """
 
+  environment_id: Optional[str] = None
+  """The execution environment ID from the interactions API.
+
+  This field is populated when an interactions-API agent (e.g. ManagedAgent)
+  provisions or reuses a sandbox environment. It is persisted on the resulting
+  Event so subsequent turns can reuse the same environment for stateful work.
+  """
+
   def get_function_calls(self) -> list[types.FunctionCall]:
     """Returns the function calls in the response."""
     func_calls = []

tests/unittests/flows/llm_flows/test_interactions_processor.py

@@ -236,11 +236,11 @@ def test_find_previous_interaction_id_returns_latest_for_agent():
       _evt("other_agent", "int_3", None),
   ]
 
-  result = interactions_processor._find_previous_interaction_id(
+  result = interactions_processor._find_previous_interaction_state(
       events, agent_name="my_agent", current_branch=None
   )
 
-  assert result == "int_2"
+  assert result[0] == "int_2"
 
 
 def test_find_previous_interaction_id_respects_branch():
@@ -249,18 +249,32 @@ def test_find_previous_interaction_id_respects_branch():
       _evt("my_agent", "int_other_branch", "branch_b"),
   ]
 
-  result = interactions_processor._find_previous_interaction_id(
+  result = interactions_processor._find_previous_interaction_state(
       events, agent_name="my_agent", current_branch="branch_a"
   )
 
-  assert result == "int_main"
+  assert result[0] == "int_main"
 
 
 def test_find_previous_interaction_id_none_when_absent():
   events = [_evt("user", None, None)]
 
-  result = interactions_processor._find_previous_interaction_id(
+  result = interactions_processor._find_previous_interaction_state(
       events, agent_name="my_agent", current_branch=None
   )
 
-  assert result is None
+  assert result[0] is None
+
+
+def test_find_previous_interaction_state_returns_both_ids():
+  events = [
+      Event(author="my_agent", interaction_id="int_1", environment_id="env_1"),
+      Event(author="user"),
+      Event(author="my_agent", interaction_id="int_2", environment_id="env_2"),
+  ]
+
+  state = interactions_processor._find_previous_interaction_state(
+      events, agent_name="my_agent", current_branch=None
+  )
+
+  assert state == ("int_2", "env_2")

tests/unittests/models/test_interactions_utils.py

@@ -2028,3 +2028,122 @@ async def test_generate_content_via_interactions_non_streaming_yields_single_res
   assert len(responses) == 1
   assert responses[0].interaction_id == 'interaction_ns'
   assert responses[0].content.parts[0].text == 'Sunny in Tokyo.'
+
+
+def _build_stream_with_environment() -> list[object]:
+  """A streamed interaction whose completed event carries an environment id."""
+  now = datetime.now(timezone.utc).isoformat()
+  created = InteractionCreatedEvent(
+      event_type='interaction.created',
+      interaction=InteractionSseEventInteraction(
+          id='interaction_env',
+          created=now,
+          updated=now,
+          status='requires_action',
+          steps=[],
+      ),
+  )
+  step_start = StepStart(
+      event_type='step.start',
+      index=0,
+      step=ModelOutputStep(type='model_output'),
+  )
+  step_delta = StepDelta(
+      event_type='step.delta',
+      index=0,
+      delta={'type': 'text', 'text': 'hi'},
+  )
+  step_stop = StepStop(event_type='step.stop', index=0)
+  completed = InteractionCompletedEvent(
+      event_type='interaction.completed',
+      interaction=InteractionSseEventInteraction(
+          id='interaction_env',
+          created=now,
+          updated=now,
+          status='completed',
+          environment_id='env_xyz',
+          steps=[
+              ModelOutputStep(
+                  type='model_output',
+                  content=[TextContent(type='text', text='hi')],
+              )
+          ],
+      ),
+  )
+  return [created, step_start, step_delta, step_stop, completed]
+
+
+def test_create_interactions_surfaces_environment_id():
+  api_client = _FakeApiClient(_build_stream_with_environment())
+
+  async def _collect():
+    out = []
+    async for r in interactions_utils._create_interactions(
+        api_client,
+        create_kwargs={'agent': 'agents/a', 'input': []},
+        stream=True,
+    ):
+      out.append(r)
+    return out
+
+  responses = asyncio.run(_collect())
+  assert responses, 'expected streamed responses'
+  # The env id arrives only on the completed event, so earlier partial
+  # responses carry no environment id.
+  assert responses[0].environment_id is None
+  assert responses[-1].environment_id == 'env_xyz'
+
+
+class _FakeNonStreamInteractions:
+  """Fake interactions resource returning a full Interaction (non-streaming)."""
+
+  def __init__(self, interaction: Interaction):
+    self._interaction = interaction
+
+  async def create(self, **_kwargs):
+    return self._interaction
+
+
+class _FakeNonStreamAio:
+  """Namespace matching the expected api_client.aio shape (non-streaming)."""
+
+  def __init__(self, interaction: Interaction):
+    self.interactions = _FakeNonStreamInteractions(interaction)
+
+
+class _FakeNonStreamApiClient:
+  """Minimal fake API client whose create() returns a full Interaction."""
+
+  def __init__(self, interaction: Interaction):
+    self.aio = _FakeNonStreamAio(interaction)
+
+
+def test_create_interactions_surfaces_environment_id_non_stream():
+  interaction = Interaction(
+      id='interaction_ns',
+      status='completed',
+      created=datetime.now(timezone.utc).isoformat(),
+      updated=datetime.now(timezone.utc).isoformat(),
+      environment_id='env_ns',
+      steps=[
+          ModelOutputStep(
+              type='model_output',
+              content=[TextContent(type='text', text='hi')],
+          )
+      ],
+  )
+  api_client = _FakeNonStreamApiClient(interaction)
+
+  async def _collect():
+    out = []
+    async for r in interactions_utils._create_interactions(
+        api_client,
+        create_kwargs={'agent': 'agents/a', 'input': []},
+        stream=False,
+    ):
+      out.append(r)
+    return out
+
+  responses = asyncio.run(_collect())
+  assert len(responses) == 1
+  assert responses[-1].environment_id == 'env_ns'

tests/unittests/models/test_llm_response.py

@@ -447,3 +447,13 @@ def test_get_function_responses_empty_when_no_content():
 def test_get_function_responses_empty_when_no_parts():
   response = LlmResponse(content=types.Content(parts=None))
   assert response.get_function_responses() == []
+
+
+def test_environment_id_defaults_to_none_and_roundtrips():
+  resp = LlmResponse()
+  assert resp.environment_id is None
+
+  resp.environment_id = 'env_abc'
+  dumped = resp.model_dump(exclude_none=True)
+  assert dumped['environment_id'] == 'env_abc'
+  assert LlmResponse.model_validate(dumped).environment_id == 'env_abc'