chore: import upstream snapshot with attribution

tests/unittests/labs/antigravity/test_antigravity_agent.py

@@ -0,0 +1,123 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for AntigravityAgent.
+
+Verifies the root-only construction constraint that keeps the agent usable only
+as a standalone root agent while the SDK supports local mode only.
+"""
+
+from __future__ import annotations
+
+from unittest.mock import AsyncMock
+from unittest.mock import MagicMock
+from unittest.mock import patch
+
+from google.adk.agents.base_agent import BaseAgent
+from google.adk.labs.antigravity import _antigravity_agent
+from google.adk.labs.antigravity._antigravity_agent import AntigravityAgent
+from google.antigravity import LocalAgentConfig
+import pytest
+
+
+def _make_config(**kwargs) -> LocalAgentConfig:
+  """Returns a minimal real LocalAgentConfig for the wrapped SDK agent."""
+  return LocalAgentConfig(system_instructions='test', **kwargs)
+
+
+def test_standalone_agent_is_allowed():
+  """An AntigravityAgent with no parent and no sub-agents constructs cleanly."""
+  agent = AntigravityAgent(name='agy', config=_make_config())
+
+  assert agent.parent_agent is None
+  assert agent.sub_agents == []
+
+
+def test_giving_sub_agents_is_rejected():
+  """Constructing with sub-agents raises a temporary root-only error."""
+  child = BaseAgent(name='child')
+
+  with pytest.raises(ValueError, match='standalone root agent'):
+    AntigravityAgent(name='agy', config=_make_config(), sub_agents=[child])
+
+
+def test_using_as_sub_agent_is_rejected():
+  """Adopting the agent under a parent raises a temporary root-only error."""
+  agy = AntigravityAgent(name='agy', config=_make_config())
+
+  with pytest.raises(ValueError, match='standalone root agent'):
+    BaseAgent(name='parent', sub_agents=[agy])
+
+
+@pytest.mark.asyncio
+async def test_run_without_save_dir_raises():
+  """Running without config.save_dir raises, since trajectories need a folder."""
+  agent = AntigravityAgent(name='agy', config=_make_config())
+
+  with pytest.raises(ValueError, match='requires config.save_dir'):
+    async for _ in agent._run_async_impl(MagicMock()):
+      pass
+
+
+@pytest.mark.asyncio
+async def test_resumed_replayed_steps_are_skipped(tmp_path):
+  """On resume, steps at or below the resume index are not re-emitted."""
+  from google.antigravity import types as sdk_types
+
+  def _step(step_index: int, text: str):
+    step = MagicMock()
+    step.step_index = step_index
+    step.source = sdk_types.StepSource.MODEL
+    step.type = sdk_types.StepType.TEXT_RESPONSE
+    step.status = sdk_types.StepStatus.DONE
+    step.is_complete_response = True
+    step.content = text
+    step.tool_calls = []
+    return step
+
+  # The harness replays steps 0-1 (prior turn) then emits step 2 (this turn).
+  async def _receive_steps():
+    yield _step(0, 'old-1')
+    yield _step(1, 'old-2')
+    yield _step(2, 'new')
+
+  conversation = MagicMock()
+  conversation.send = AsyncMock()
+  conversation.receive_steps = _receive_steps
+  active_agent = MagicMock()
+  active_agent.conversation = conversation
+  active_agent.conversation_id = 'sess_456_agy'
+  active_agent.__aenter__ = AsyncMock(return_value=active_agent)
+  active_agent.__aexit__ = AsyncMock(return_value=None)
+
+  # A prior trajectory + resume index in save_dir triggers resume at index 1.
+  save_dir = tmp_path
+  (save_dir / 'traj-sess_456_agy').write_bytes(b'data')
+  (save_dir / 'traj-sess_456_agy.resume').write_text('1')
+  agent = AntigravityAgent(
+      name='agy', config=_make_config(save_dir=str(save_dir))
+  )
+
+  ctx = MagicMock()
+  ctx.invocation_id = 'inv_1'
+  ctx.branch = 'main'
+  ctx.session.id = 'sess_456'
+  ctx.user_content = None
+  ctx.run_config = None
+
+  with patch.object(_antigravity_agent, 'Agent', return_value=active_agent):
+    events = [event async for event in agent._run_async_impl(ctx)]
+
+  texts = [e.content.parts[0].text for e in events]
+  assert texts == ['new']

tests/unittests/labs/antigravity/test_event_converter.py

@@ -0,0 +1,226 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for the Antigravity step-to-event converter.
+
+Verifies that model text, function calls, and function responses map to the
+expected ADK events, and that repeated steps are deduplicated.
+"""
+
+from __future__ import annotations
+
+from unittest.mock import MagicMock
+
+from google.adk.labs.antigravity import _event_converter
+from google.antigravity import types as sdk_types
+
+
+def _make_ctx() -> MagicMock:
+  ctx = MagicMock()
+  ctx.invocation_id = 'inv_1'
+  ctx.branch = 'main'
+  return ctx
+
+
+def _convert(step, *, streaming=False):
+  return _event_converter.convert_step_to_events(
+      step,
+      ctx=_make_ctx(),
+      author='agy',
+      seen_tool_calls=set(),
+      seen_tool_results=set(),
+      streaming=streaming,
+  )
+
+
+def test_completed_model_text_maps_to_one_model_text_event():
+  """A completed model text response becomes a single model text event."""
+  step = sdk_types.Step(
+      step_index=0,
+      type=sdk_types.StepType.TEXT_RESPONSE,
+      source=sdk_types.StepSource.MODEL,
+      content='hello there',
+      is_complete_response=True,
+  )
+
+  events = _convert(step)
+
+  assert len(events) == 1
+  assert events[0].author == 'agy'
+  assert events[0].content.role == 'model'
+  assert events[0].content.parts[0].text == 'hello there'
+
+
+def test_partial_model_text_produces_no_event():
+  """A streaming partial text step (cumulative snapshot) yields nothing."""
+  step = sdk_types.Step(
+      step_index=0,
+      type=sdk_types.StepType.TEXT_RESPONSE,
+      source=sdk_types.StepSource.MODEL,
+      content='hello',
+      content_delta='hello',
+      is_complete_response=None,
+  )
+
+  assert _convert(step) == []
+
+
+def test_function_call_maps_to_function_call_event():
+  """A model tool-call step becomes a model function-call event."""
+  step = sdk_types.Step(
+      step_index=1,
+      type=sdk_types.StepType.TOOL_CALL,
+      source=sdk_types.StepSource.MODEL,
+      tool_calls=[
+          sdk_types.ToolCall(name='view_file', args={'path': '/x'}, id='c1')
+      ],
+  )
+
+  events = _convert(step)
+
+  assert len(events) == 1
+  fc = events[0].content.parts[0].function_call
+  assert events[0].author == 'agy'
+  assert fc.name == 'view_file'
+  assert fc.id == 'c1'
+  assert fc.args == {'path': '/x'}
+
+
+def test_function_response_maps_to_function_response_event():
+  """A completed tool-execution step becomes a function-response event."""
+  step = sdk_types.Step(
+      step_index=2,
+      type=sdk_types.StepType.TOOL_CALL,
+      source=sdk_types.StepSource.SYSTEM,
+      status=sdk_types.StepStatus.DONE,
+      content='file contents',
+      tool_calls=[sdk_types.ToolCall(name='view_file', args={}, id='c1')],
+  )
+
+  events = _convert(step)
+
+  assert len(events) == 1
+  fr = events[0].content.parts[0].function_response
+  assert events[0].author == 'view_file'
+  assert events[0].content.role == 'user'
+  assert fr.name == 'view_file'
+  assert fr.id == 'c1'
+  assert fr.response == {'result': 'file contents'}
+
+
+def test_errored_tool_step_maps_error_response():
+  """A failed tool-execution step reports the error in the response payload."""
+  step = sdk_types.Step(
+      step_index=3,
+      type=sdk_types.StepType.TOOL_CALL,
+      source=sdk_types.StepSource.SYSTEM,
+      status=sdk_types.StepStatus.ERROR,
+      error='permission denied',
+      tool_calls=[sdk_types.ToolCall(name='run_command', args={}, id='c2')],
+  )
+
+  events = _convert(step)
+
+  assert events[0].content.parts[0].function_response.response == {
+      'error': 'permission denied'
+  }
+
+
+def test_duplicate_tool_call_emitted_once():
+  """The same tool call repeated across steps is emitted only once."""
+  call = sdk_types.ToolCall(name='view_file', args={}, id='c1')
+  step = sdk_types.Step(
+      step_index=1,
+      type=sdk_types.StepType.TOOL_CALL,
+      source=sdk_types.StepSource.MODEL,
+      tool_calls=[call],
+  )
+  ctx = _make_ctx()
+  seen: set[str] = set()
+
+  first = _event_converter.convert_step_to_events(
+      step, ctx=ctx, author='agy', seen_tool_calls=seen, seen_tool_results=set()
+  )
+  second = _event_converter.convert_step_to_events(
+      step, ctx=ctx, author='agy', seen_tool_calls=seen, seen_tool_results=set()
+  )
+
+  assert len(first) == 1
+  assert second == []
+
+
+def test_incomplete_text_step_produces_no_final_event():
+  """A non-final text step yields nothing in non-streaming mode."""
+  step = sdk_types.Step(
+      step_index=0,
+      type=sdk_types.StepType.TEXT_RESPONSE,
+      source=sdk_types.StepSource.MODEL,
+      thinking='reasoning...',
+      content='',
+  )
+
+  assert _convert(step) == []
+
+
+def test_streaming_emits_partial_thinking_then_text_deltas():
+  """In SSE mode a step's thinking and text deltas become partial events."""
+  step = sdk_types.Step(
+      step_index=0,
+      type=sdk_types.StepType.TEXT_RESPONSE,
+      source=sdk_types.StepSource.MODEL,
+      thinking_delta='thinking...',
+      content_delta='hello',
+  )
+
+  events = _convert(step, streaming=True)
+
+  assert len(events) == 2
+  assert events[0].partial is True
+  assert events[0].content.parts[0].thought is True
+  assert events[0].content.parts[0].text == 'thinking...'
+  assert events[1].partial is True
+  assert events[1].content.parts[0].text == 'hello'
+
+
+def test_non_streaming_omits_partial_deltas():
+  """Without SSE mode, delta-only steps yield no events."""
+  step = sdk_types.Step(
+      step_index=0,
+      type=sdk_types.StepType.TEXT_RESPONSE,
+      source=sdk_types.StepSource.MODEL,
+      thinking_delta='thinking...',
+      content_delta='hello',
+  )
+
+  assert _convert(step, streaming=False) == []
+
+
+def test_streaming_completed_step_emits_partial_then_final():
+  """A completed step in SSE mode emits the partial delta then the final text."""
+  step = sdk_types.Step(
+      step_index=1,
+      type=sdk_types.StepType.TEXT_RESPONSE,
+      source=sdk_types.StepSource.MODEL,
+      content_delta=' world',
+      content='hello world',
+      is_complete_response=True,
+  )
+
+  events = _convert(step, streaming=True)
+
+  assert len(events) == 2
+  assert events[0].partial is True
+  assert events[0].content.parts[0].text == ' world'
+  assert events[1].partial in (False, None)
+  assert events[1].content.parts[0].text == 'hello world'

tests/unittests/labs/antigravity/test_trajectory_files.py

@@ -0,0 +1,79 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for Antigravity trajectory resumption bookkeeping in save_dir.
+
+Verifies trajectory detection, resume step index persistence, and renaming the
+harness's randomly-named trajectory to a deterministic name.
+"""
+
+from __future__ import annotations
+
+from google.adk.labs.antigravity import _trajectory_files
+
+
+def test_has_trajectory_false_when_absent(tmp_path):
+  """No trajectory file means no prior conversation to resume."""
+  assert not _trajectory_files.has_trajectory(str(tmp_path), 'sess_agy')
+
+
+def test_has_trajectory_true_when_present(tmp_path):
+  """An existing traj file is detected for the conversation."""
+  (tmp_path / 'traj-sess_agy').write_bytes(b'data')
+
+  assert _trajectory_files.has_trajectory(str(tmp_path), 'sess_agy')
+
+
+def test_load_resume_step_index_minus_one_when_absent(tmp_path):
+  """Missing resume step index reads as -1 (fresh)."""
+  assert (
+      _trajectory_files.load_resume_step_index(str(tmp_path), 'sess_agy') == -1
+  )
+
+
+def test_resume_step_index_round_trips(tmp_path):
+  """A saved resume step index reads back as the same value."""
+  _trajectory_files.save_resume_step_index(str(tmp_path), 'sess_agy', 12)
+
+  assert (
+      _trajectory_files.load_resume_step_index(str(tmp_path), 'sess_agy') == 12
+  )
+
+
+def test_load_resume_step_index_minus_one_when_corrupt(tmp_path):
+  """A non-integer resume step index is treated as fresh."""
+  (tmp_path / 'traj-sess_agy.resume').write_text('not-an-int')
+
+  assert (
+      _trajectory_files.load_resume_step_index(str(tmp_path), 'sess_agy') == -1
+  )
+
+
+def test_rename_trajectory_to_conversation_id(tmp_path):
+  """The harness's random trajectory is renamed to the deterministic name."""
+  (tmp_path / 'traj-random123').write_bytes(b'data')
+
+  _trajectory_files.rename_trajectory(str(tmp_path), 'sess_agy', 'random123')
+
+  assert not (tmp_path / 'traj-random123').exists()
+  assert (tmp_path / 'traj-sess_agy').read_bytes() == b'data'
+
+
+def test_rename_trajectory_noop_when_already_named(tmp_path):
+  """Renaming is a no-op when the harness id already matches."""
+  (tmp_path / 'traj-sess_agy').write_bytes(b'data')
+
+  _trajectory_files.rename_trajectory(str(tmp_path), 'sess_agy', 'sess_agy')
+
+  assert (tmp_path / 'traj-sess_agy').read_bytes() == b'data'