chore: import upstream snapshot with attribution

tests/unittests/runners/test_runner_rewind.py

@@ -0,0 +1,363 @@
+# 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 runner.rewind_async."""
+
+from typing import Any
+from typing import Optional
+from typing import Union
+
+from google.adk.agents.base_agent import BaseAgent
+from google.adk.artifacts.base_artifact_service import ensure_part
+from google.adk.artifacts.in_memory_artifact_service import InMemoryArtifactService
+from google.adk.events.event import Event
+from google.adk.events.event import EventActions
+from google.adk.runners import Runner
+from google.adk.sessions.in_memory_session_service import InMemorySessionService
+from google.genai import types
+import pytest
+
+
+class _NoFileDataArtifactService(InMemoryArtifactService):
+  """Artifact service that rejects file_data parts, like GCS/File services."""
+
+  async def save_artifact(
+      self,
+      *,
+      app_name: str,
+      user_id: str,
+      filename: str,
+      artifact: Union[types.Part, dict[str, Any]],
+      session_id: Optional[str] = None,
+      custom_metadata: Optional[dict[str, Any]] = None,
+  ) -> int:
+    artifact = ensure_part(artifact)
+    if artifact.file_data:
+      raise NotImplementedError(
+          "Saving artifact with file_data is not supported."
+      )
+    return await super().save_artifact(
+        app_name=app_name,
+        user_id=user_id,
+        filename=filename,
+        artifact=artifact,
+        session_id=session_id,
+        custom_metadata=custom_metadata,
+    )
+
+
+class TestRunnerRewind:
+  """Tests for runner.rewind_async."""
+
+  runner: Runner
+
+  def setup_method(self):
+    """Set up test fixtures."""
+    root_agent = BaseAgent(name="test_agent")
+    session_service = InMemorySessionService()
+    artifact_service = InMemoryArtifactService()
+    self.runner = Runner(
+        app_name="test_app",
+        agent=root_agent,
+        session_service=session_service,
+        artifact_service=artifact_service,
+    )
+
+  @pytest.mark.asyncio
+  async def test_rewind_async_with_state_and_artifacts(self):
+    """Tests rewind_async rewinds state and artifacts."""
+    runner = self.runner
+    user_id = "test_user"
+    session_id = "test_session"
+
+    # 1. Setup session and initial artifacts
+    session = await runner.session_service.create_session(
+        app_name=runner.app_name, user_id=user_id, session_id=session_id
+    )
+
+    # invocation1
+    await runner.artifact_service.save_artifact(
+        app_name=runner.app_name,
+        user_id=user_id,
+        session_id=session_id,
+        filename="f1",
+        artifact=types.Part.from_text(text="f1v0"),
+    )
+    event1 = Event(
+        invocation_id="invocation1",
+        author="agent",
+        content=types.Content(parts=[types.Part.from_text(text="event1")]),
+        actions=EventActions(
+            state_delta={"k1": "v1"}, artifact_delta={"f1": 0}
+        ),
+    )
+    await runner.session_service.append_event(session=session, event=event1)
+
+    # invocation2
+    await runner.artifact_service.save_artifact(
+        app_name=runner.app_name,
+        user_id=user_id,
+        session_id=session_id,
+        filename="f1",
+        artifact=types.Part.from_text(text="f1v1"),
+    )
+    await runner.artifact_service.save_artifact(
+        app_name=runner.app_name,
+        user_id=user_id,
+        session_id=session_id,
+        filename="f2",
+        artifact=types.Part.from_text(text="f2v0"),
+    )
+    event2 = Event(
+        invocation_id="invocation2",
+        author="agent",
+        content=types.Content(parts=[types.Part.from_text(text="event2")]),
+        actions=EventActions(
+            state_delta={"k1": "v2", "k2": "v2"},
+            artifact_delta={"f1": 1, "f2": 0},
+        ),
+    )
+    await runner.session_service.append_event(session=session, event=event2)
+
+    # invocation3
+    event3 = Event(
+        invocation_id="invocation3",
+        author="agent",
+        content=types.Content(parts=[types.Part.from_text(text="event3")]),
+        actions=EventActions(state_delta={"k2": "v3"}),
+    )
+    await runner.session_service.append_event(session=session, event=event3)
+
+    session = await runner.session_service.get_session(
+        app_name=runner.app_name, user_id=user_id, session_id=session_id
+    )
+    assert session.state == {"k1": "v2", "k2": "v3"}
+    assert await runner.artifact_service.load_artifact(
+        app_name=runner.app_name,
+        user_id=user_id,
+        session_id=session_id,
+        filename="f1",
+    ) == types.Part.from_text(text="f1v1")
+    assert await runner.artifact_service.load_artifact(
+        app_name=runner.app_name,
+        user_id=user_id,
+        session_id=session_id,
+        filename="f2",
+    ) == types.Part.from_text(text="f2v0")
+
+    # 2. Rewind before invocation2
+    await runner.rewind_async(
+        user_id=user_id,
+        session_id=session_id,
+        rewind_before_invocation_id="invocation2",
+    )
+
+    # 3. Verify state and artifacts are rewound
+    session = await runner.session_service.get_session(
+        app_name=runner.app_name, user_id=user_id, session_id=session_id
+    )
+    # After rewind before invocation2, only event1 state delta should apply.
+    assert session.state["k1"] == "v1"
+    assert not session.state["k2"]
+    # f1 should be rewound to v0
+    assert await runner.artifact_service.load_artifact(
+        app_name=runner.app_name,
+        user_id=user_id,
+        session_id=session_id,
+        filename="f1",
+    ) == types.Part.from_text(text="f1v0")
+    # f2 should not exist
+    assert (
+        await runner.artifact_service.load_artifact(
+            app_name=runner.app_name,
+            user_id=user_id,
+            session_id=session_id,
+            filename="f2",
+        )
+        is None
+    )
+
+  @pytest.mark.asyncio
+  async def test_rewind_async_not_first_invocation(self):
+    """Tests rewind_async rewinds state and artifacts to invocation2."""
+    runner = self.runner
+    user_id = "test_user"
+    session_id = "test_session"
+
+    # 1. Setup session and initial artifacts
+    session = await runner.session_service.create_session(
+        app_name=runner.app_name, user_id=user_id, session_id=session_id
+    )
+    # invocation1
+    await runner.artifact_service.save_artifact(
+        app_name=runner.app_name,
+        user_id=user_id,
+        session_id=session_id,
+        filename="f1",
+        artifact=types.Part.from_text(text="f1v0"),
+    )
+    event1 = Event(
+        invocation_id="invocation1",
+        author="agent",
+        content=types.Content(parts=[types.Part.from_text(text="event1")]),
+        actions=EventActions(
+            state_delta={"k1": "v1"}, artifact_delta={"f1": 0}
+        ),
+    )
+    await runner.session_service.append_event(session=session, event=event1)
+
+    # invocation2
+    await runner.artifact_service.save_artifact(
+        app_name=runner.app_name,
+        user_id=user_id,
+        session_id=session_id,
+        filename="f1",
+        artifact=types.Part.from_text(text="f1v1"),
+    )
+    await runner.artifact_service.save_artifact(
+        app_name=runner.app_name,
+        user_id=user_id,
+        session_id=session_id,
+        filename="f2",
+        artifact=types.Part.from_text(text="f2v0"),
+    )
+    event2 = Event(
+        invocation_id="invocation2",
+        author="agent",
+        content=types.Content(parts=[types.Part.from_text(text="event2")]),
+        actions=EventActions(
+            state_delta={"k1": "v2", "k2": "v2"},
+            artifact_delta={"f1": 1, "f2": 0},
+        ),
+    )
+    await runner.session_service.append_event(session=session, event=event2)
+
+    # invocation3
+    event3 = Event(
+        invocation_id="invocation3",
+        author="agent",
+        content=types.Content(parts=[types.Part.from_text(text="event3")]),
+        actions=EventActions(state_delta={"k2": "v3"}),
+    )
+    await runner.session_service.append_event(session=session, event=event3)
+
+    # 2. Rewind before invocation3
+    await runner.rewind_async(
+        user_id=user_id,
+        session_id=session_id,
+        rewind_before_invocation_id="invocation3",
+    )
+
+    # 3. Verify state and artifacts are rewound
+    session = await runner.session_service.get_session(
+        app_name=runner.app_name, user_id=user_id, session_id=session_id
+    )
+    # After rewind before invocation3, event1 and event2 state deltas should apply.
+    assert session.state == {"k1": "v2", "k2": "v2"}
+    # f1 should be v1
+    assert await runner.artifact_service.load_artifact(
+        app_name=runner.app_name,
+        user_id=user_id,
+        session_id=session_id,
+        filename="f1",
+    ) == types.Part.from_text(text="f1v1")
+    # f2 should be v0
+    assert await runner.artifact_service.load_artifact(
+        app_name=runner.app_name,
+        user_id=user_id,
+        session_id=session_id,
+        filename="f2",
+    ) == types.Part.from_text(text="f2v0")
+
+
+class TestRunnerRewindNoFileData:
+  """Tests that rewind works with artifact services that reject file_data."""
+
+  @pytest.mark.asyncio
+  async def test_rewind_uses_load_artifact_not_file_data(self):
+    """Rewind must not construct file_data parts for artifact restoration.
+
+    GCS and File artifact services reject file_data parts. The runner
+    should use load_artifact to get inline_data instead.
+    """
+    root_agent = BaseAgent(name="test_agent")
+    session_service = InMemorySessionService()
+    artifact_service = _NoFileDataArtifactService()
+    runner = Runner(
+        app_name="test_app",
+        agent=root_agent,
+        session_service=session_service,
+        artifact_service=artifact_service,
+    )
+    user_id = "test_user"
+    session_id = "test_session"
+
+    session = await runner.session_service.create_session(
+        app_name=runner.app_name, user_id=user_id, session_id=session_id
+    )
+
+    # invocation1: create artifact f1 v0
+    await runner.artifact_service.save_artifact(
+        app_name=runner.app_name,
+        user_id=user_id,
+        session_id=session_id,
+        filename="f1",
+        artifact=types.Part.from_text(text="f1v0"),
+    )
+    event1 = Event(
+        invocation_id="invocation1",
+        author="agent",
+        content=types.Content(parts=[types.Part.from_text(text="e1")]),
+        actions=EventActions(
+            state_delta={"k1": "v1"}, artifact_delta={"f1": 0}
+        ),
+    )
+    await runner.session_service.append_event(session=session, event=event1)
+
+    # invocation2: update artifact f1 to v1
+    await runner.artifact_service.save_artifact(
+        app_name=runner.app_name,
+        user_id=user_id,
+        session_id=session_id,
+        filename="f1",
+        artifact=types.Part.from_text(text="f1v1"),
+    )
+    event2 = Event(
+        invocation_id="invocation2",
+        author="agent",
+        content=types.Content(parts=[types.Part.from_text(text="e2")]),
+        actions=EventActions(artifact_delta={"f1": 1}),
+    )
+    await runner.session_service.append_event(session=session, event=event2)
+
+    session = await runner.session_service.get_session(
+        app_name=runner.app_name, user_id=user_id, session_id=session_id
+    )
+
+    # Rewind before invocation2 — this would raise NotImplementedError
+    # with the old code that constructed file_data parts.
+    await runner.rewind_async(
+        user_id=user_id,
+        session_id=session_id,
+        rewind_before_invocation_id="invocation2",
+    )
+
+    # f1 should be restored to v0 content
+    restored = await runner.artifact_service.load_artifact(
+        app_name=runner.app_name,
+        user_id=user_id,
+        session_id=session_id,
+        filename="f1",
+    )
+    assert restored == types.Part.from_text(text="f1v0")