feat: Add rewind_async to support rewinding the session to before a previous invocation

PiperOrigin-RevId: 820552460

contributing/samples/rewind_session/__init__.py

@@ -0,0 +1,15 @@
+# Copyright 2025 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.
+
+from . import agent

contributing/samples/rewind_session/agent.py

@@ -0,0 +1,71 @@
+# Copyright 2025 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.
+
+from google.adk import Agent
+from google.adk.tools.tool_context import ToolContext
+from google.genai import types
+
+
+async def update_state(tool_context: ToolContext, key: str, value: str) -> dict:
+  """Updates a state value."""
+  tool_context.state[key] = value
+  return {"status": f"Updated state '{key}' to '{value}'"}
+
+
+async def load_state(tool_context: ToolContext, key: str) -> dict:
+  """Loads a state value."""
+  return {key: tool_context.state.get(key)}
+
+
+async def save_artifact(
+    tool_context: ToolContext, filename: str, content: str
+) -> dict:
+  """Saves an artifact with the given filename and content."""
+  artifact_bytes = content.encode("utf-8")
+  artifact_part = types.Part(
+      inline_data=types.Blob(mime_type="text/plain", data=artifact_bytes)
+  )
+  version = await tool_context.save_artifact(filename, artifact_part)
+  return {"status": "success", "filename": filename, "version": version}
+
+
+async def load_artifact(tool_context: ToolContext, filename: str) -> dict:
+  """Loads an artifact with the given filename."""
+  artifact = await tool_context.load_artifact(filename)
+  if not artifact:
+    return {"error": f"Artifact '{filename}' not found"}
+  content = artifact.inline_data.data.decode("utf-8")
+  return {"filename": filename, "content": content}
+
+
+# Create the agent
+root_agent = Agent(
+    name="state_agent",
+    model="gemini-2.0-flash",
+    instruction="""You are an agent that manages state and artifacts.
+
+    You can:
+    - Update state value
+    - Load state value
+    - Save artifact
+    - Load artifact
+
+    Use the appropriate tool based on what the user asks for.""",
+    tools=[
+        update_state,
+        load_state,
+        save_artifact,
+        load_artifact,
+    ],
+)

src/google/adk/artifacts/artifact_util.py

@@ -0,0 +1,116 @@
+# Copyright 2025 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.
+"""Utility functions for handling artifact URIs."""
+
+from __future__ import annotations
+
+import re
+from typing import NamedTuple
+from typing import Optional
+
+from google.genai import types
+
+
+class ParsedArtifactUri(NamedTuple):
+  """The result of parsing an artifact URI."""
+
+  app_name: str
+  user_id: str
+  session_id: Optional[str]
+  filename: str
+  version: int
+
+
+_SESSION_SCOPED_ARTIFACT_URI_RE = re.compile(
+    r"artifact://apps/([^/]+)/users/([^/]+)/sessions/([^/]+)/artifacts/([^/]+)/versions/(\d+)"
+)
+_USER_SCOPED_ARTIFACT_URI_RE = re.compile(
+    r"artifact://apps/([^/]+)/users/([^/]+)/artifacts/([^/]+)/versions/(\d+)"
+)
+
+
+def parse_artifact_uri(uri: str) -> Optional[ParsedArtifactUri]:
+  """Parses an artifact URI.
+
+  Args:
+      uri: The artifact URI to parse.
+
+  Returns:
+      A ParsedArtifactUri if parsing is successful, None otherwise.
+  """
+  if not uri or not uri.startswith("artifact://"):
+    return None
+
+  match = _SESSION_SCOPED_ARTIFACT_URI_RE.match(uri)
+  if match:
+    return ParsedArtifactUri(
+        app_name=match.group(1),
+        user_id=match.group(2),
+        session_id=match.group(3),
+        filename=match.group(4),
+        version=int(match.group(5)),
+    )
+
+  match = _USER_SCOPED_ARTIFACT_URI_RE.match(uri)
+  if match:
+    return ParsedArtifactUri(
+        app_name=match.group(1),
+        user_id=match.group(2),
+        session_id=None,
+        filename=match.group(3),
+        version=int(match.group(4)),
+    )
+
+  return None
+
+
+def get_artifact_uri(
+    app_name: str,
+    user_id: str,
+    filename: str,
+    version: int,
+    session_id: Optional[str] = None,
+) -> str:
+  """Constructs an artifact URI.
+
+  Args:
+      app_name: The name of the application.
+      user_id: The ID of the user.
+      filename: The name of the artifact file.
+      version: The version of the artifact.
+      session_id: The ID of the session.
+
+  Returns:
+      The constructed artifact URI.
+  """
+  if session_id:
+    return f"artifact://apps/{app_name}/users/{user_id}/sessions/{session_id}/artifacts/{filename}/versions/{version}"
+  else:
+    return f"artifact://apps/{app_name}/users/{user_id}/artifacts/{filename}/versions/{version}"
+
+
+def is_artifact_ref(artifact: types.Part) -> bool:
+  """Checks if an artifact part is an artifact reference.
+
+  Args:
+      artifact: The artifact part to check.
+
+  Returns:
+      True if the artifact part is an artifact reference, False otherwise.
+  """
+  return bool(
+      artifact.file_data
+      and artifact.file_data.file_uri
+      and artifact.file_data.file_uri.startswith("artifact://")
+  )

src/google/adk/artifacts/gcs_artifact_service.py

@@ -212,6 +212,11 @@ class GcsArtifactService(BaseArtifactService):
       blob.upload_from_string(
           data=artifact.text,
       )
+    elif artifact.file_data:
+      raise NotImplementedError(
+          "Saving artifact with file_data is not supported yet in"
+          " GcsArtifactService."
+      )
     else:
       raise ValueError("Artifact must have either inline_data or text.")
 

src/google/adk/artifacts/in_memory_artifact_service.py

@@ -18,6 +18,7 @@ import logging
 from typing import Any
 from typing import Optional
 
+from google.adk.artifacts import artifact_util
 from google.genai import types
 from pydantic import BaseModel
 from pydantic import Field
@@ -122,7 +123,15 @@ class InMemoryArtifactService(BaseArtifactService, BaseModel):
     elif artifact.text is not None:
       artifact_version.mime_type = "text/plain"
     elif artifact.file_data is not None:
-      artifact_version.mime_type = artifact.file_data.mime_type
+      if artifact_util.is_artifact_ref(artifact):
+        if not artifact_util.parse_artifact_uri(artifact.file_data.file_uri):
+          raise ValueError(
+              f"Invalid artifact reference URI: {artifact.file_data.file_uri}"
+          )
+        # If it's a valid artifact URI, we store the artifact part as-is.
+        # And we don't know the mime type until we load it.
+      else:
+        artifact_version.mime_type = artifact.file_data.mime_type
     else:
       raise ValueError("Not supported artifact type.")
 
@@ -147,11 +156,42 @@ class InMemoryArtifactService(BaseArtifactService, BaseModel):
       return None
     if version is None:
       version = -1
+
     try:
-      return versions[version].data
+      artifact_entry = versions[version]
     except IndexError:
       return None
 
+    if artifact_entry is None:
+      return None
+
+    # Resolve artifact reference if needed.
+    artifact_data = artifact_entry.data
+    if artifact_util.is_artifact_ref(artifact_data):
+      parsed_uri = artifact_util.parse_artifact_uri(
+          artifact_data.file_data.file_uri
+      )
+      if not parsed_uri:
+        raise ValueError(
+            "Invalid artifact reference URI:"
+            f" {artifact_data.file_data.file_uri}"
+        )
+      return await self.load_artifact(
+          app_name=parsed_uri.app_name,
+          user_id=parsed_uri.user_id,
+          filename=parsed_uri.filename,
+          session_id=parsed_uri.session_id,
+          version=parsed_uri.version,
+      )
+
+    if (
+        artifact_data == types.Part()
+        or artifact_data == types.Part(text="")
+        or (artifact_data.inline_data and not artifact_data.inline_data.data)
+    ):
+      return None
+    return artifact_data
+
   @override
   async def list_artifact_keys(
       self, *, app_name: str, user_id: str, session_id: Optional[str] = None

src/google/adk/events/event_actions.py

@@ -104,3 +104,6 @@ class EventActions(BaseModel):
 
   agent_state: Optional[dict[str, Any]] = None
   """The agent state at the current event."""
+
+  rewind_before_invocation_id: Optional[str] = None
+  """The invocation id to rewind to. This is only set for rewind event."""

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

@@ -310,11 +310,30 @@ def _get_contents(
   accumulated_input_transcription = ''
   accumulated_output_transcription = ''
 
+  # Filter out events that are annulled by a rewind.
+  # By iterating backward, when a rewind event is found, we skip all events
+  # from that point back to the `rewind_before_invocation_id`, thus removing
+  # them from the history used for the LLM request.
+  rewind_filtered_events = []
+  i = len(events) - 1
+  while i >= 0:
+    event = events[i]
+    if event.actions and event.actions.rewind_before_invocation_id:
+      rewind_invocation_id = event.actions.rewind_before_invocation_id
+      for j in range(0, i, 1):
+        if events[j].invocation_id == rewind_invocation_id:
+          i = j
+          break
+    else:
+      rewind_filtered_events.append(event)
+    i -= 1
+  rewind_filtered_events.reverse()
+
   # Parse the events, leaving the contents and the function calls and
   # responses from the current agent.
   raw_filtered_events = []
   has_compaction_events = False
-  for event in events:
+  for event in rewind_filtered_events:
     if _contains_empty_content(event):
       continue
     if not _is_event_belongs_to_branch(current_branch, event):

src/google/adk/runners.py

@@ -28,6 +28,7 @@ from typing import Optional
 import warnings
 
 from google.adk.apps.compaction import _run_compaction_for_sliding_window
+from google.adk.artifacts import artifact_util
 from google.genai import types
 
 from .agents.active_streaming_tool import ActiveStreamingTool
@@ -427,6 +428,146 @@ class Runner:
       async for event in agen:
         yield event
 
+  async def rewind_async(
+      self,
+      *,
+      user_id: str,
+      session_id: str,
+      rewind_before_invocation_id: str,
+  ) -> None:
+    """Rewinds the session to before the specified invocation."""
+    session = await self.session_service.get_session(
+        app_name=self.app_name, user_id=user_id, session_id=session_id
+    )
+    if not session:
+      raise ValueError(f'Session not found: {session_id}')
+
+    rewind_event_index = -1
+    for i, event in enumerate(session.events):
+      if event.invocation_id == rewind_before_invocation_id:
+        rewind_event_index = i
+        break
+
+    if rewind_event_index == -1:
+      raise ValueError(
+          f'Invocation ID not found: {rewind_before_invocation_id}'
+      )
+
+    # Compute state delta to reverse changes
+    state_delta = await self._compute_state_delta_for_rewind(
+        session, rewind_event_index
+    )
+
+    # Compute artifact delta to reverse changes
+    artifact_delta = await self._compute_artifact_delta_for_rewind(
+        session, rewind_event_index
+    )
+
+    # Create rewind event
+    rewind_event = Event(
+        invocation_id=new_invocation_context_id(),
+        author='user',
+        actions=EventActions(
+            rewind_before_invocation_id=rewind_before_invocation_id,
+            state_delta=state_delta,
+            artifact_delta=artifact_delta,
+        ),
+    )
+
+    logger.info('Rewinding session to invocation: %s', rewind_event)
+
+    await self.session_service.append_event(session=session, event=rewind_event)
+
+  async def _compute_state_delta_for_rewind(
+      self, session: Session, rewind_event_index: int
+  ) -> dict[str, Any]:
+    """Computes the state delta to reverse changes."""
+    state_at_rewind_point: dict[str, Any] = {}
+    for i in range(rewind_event_index):
+      if session.events[i].actions.state_delta:
+        for k, v in session.events[i].actions.state_delta.items():
+          if k.startswith('app:') or k.startswith('user:'):
+            continue
+          if v is None:
+            state_at_rewind_point.pop(k, None)
+          else:
+            state_at_rewind_point[k] = v
+
+    current_state = session.state
+    rewind_state_delta = {}
+
+    # 1. Add/update keys in rewind_state_delta to match state_at_rewind_point.
+    for key, value_at_rewind in state_at_rewind_point.items():
+      if key not in current_state or current_state[key] != value_at_rewind:
+        rewind_state_delta[key] = value_at_rewind
+
+    # 2. Set keys to None in rewind_state_delta if they are in current_state
+    #    but not in state_at_rewind_point. These keys were added after the
+    #    rewind point and need to be removed.
+    for key in current_state:
+      if key.startswith('app:') or key.startswith('user:'):
+        continue
+      if key not in state_at_rewind_point:
+        rewind_state_delta[key] = None
+
+    return rewind_state_delta
+
+  async def _compute_artifact_delta_for_rewind(
+      self, session: Session, rewind_event_index: int
+  ) -> dict[str, int]:
+    """Computes the artifact delta to reverse changes."""
+    if not self.artifact_service:
+      return {}
+
+    versions_at_rewind_point: dict[str, int] = {}
+    for i in range(rewind_event_index):
+      event = session.events[i]
+      if event.actions.artifact_delta:
+        versions_at_rewind_point.update(event.actions.artifact_delta)
+
+    current_versions: dict[str, int] = {}
+    for event in session.events:
+      if event.actions.artifact_delta:
+        current_versions.update(event.actions.artifact_delta)
+
+    rewind_artifact_delta = {}
+    for filename, vn in current_versions.items():
+      if filename.startswith('user:'):
+        # User artifacts are not restored on rewind.
+        continue
+      vt = versions_at_rewind_point.get(filename)
+      if vt == vn:
+        continue
+
+      rewind_artifact_delta[filename] = vn + 1
+      if vt is None:
+        # Artifact did not exist at rewind point. Mark it as inaccessible.
+        artifact = types.Part(
+            inline_data=types.Blob(
+                mime_type='application/octet-stream', data=b''
+            )
+        )
+      else:
+        # Artifact version changed after rewind point. Restore to version at
+        # rewind point.
+        artifact_uri = artifact_util.get_artifact_uri(
+            app_name=self.app_name,
+            user_id=session.user_id,
+            session_id=session.id,
+            filename=filename,
+            version=vt,
+        )
+        artifact = types.Part(file_data=types.FileData(file_uri=artifact_uri))
+      await self.artifact_service.save_artifact(
+          app_name=self.app_name,
+          user_id=session.user_id,
+          session_id=session.id,
+          filename=filename,
+          artifact=artifact,
+      )
+
+    return rewind_artifact_delta
+
   async def _run_compaction_default(self, session: Session):
     """Runs compaction for other types of compactors.
 
@@ -1083,7 +1224,7 @@ class InMemoryRunner(Runner):
       self,
       agent: Optional[BaseAgent] = None,
       *,
-      app_name: Optional[str] = 'InMemoryRunner',
+      app_name: Optional[str] = None,
       plugins: Optional[list[BasePlugin]] = None,
       app: Optional[App] = None,
   ):
@@ -1094,6 +1235,8 @@ class InMemoryRunner(Runner):
         app_name: The application name of the runner. Defaults to
           'InMemoryRunner'.
     """
+    if app is None and app_name is None:
+      app_name = 'InMemoryRunner'
     super().__init__(
         app_name=app_name,
         agent=agent,

tests/unittests/artifacts/test_artifact_util.py

@@ -0,0 +1,109 @@
+# Copyright 2025 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 artifact_util."""
+
+from google.adk.artifacts import artifact_util
+from google.genai import types
+import pytest
+
+
+def test_parse_session_scoped_artifact_uri():
+  """Tests parsing a valid session-scoped artifact URI."""
+  uri = "artifact://apps/app1/users/user1/sessions/session1/artifacts/file1/versions/123"
+  parsed = artifact_util.parse_artifact_uri(uri)
+  assert parsed is not None
+  assert parsed.app_name == "app1"
+  assert parsed.user_id == "user1"
+  assert parsed.session_id == "session1"
+  assert parsed.filename == "file1"
+  assert parsed.version == 123
+
+
+def test_parse_user_scoped_artifact_uri():
+  """Tests parsing a valid user-scoped artifact URI."""
+  uri = "artifact://apps/app2/users/user2/artifacts/file2/versions/456"
+  parsed = artifact_util.parse_artifact_uri(uri)
+  assert parsed is not None
+  assert parsed.app_name == "app2"
+  assert parsed.user_id == "user2"
+  assert parsed.session_id is None
+  assert parsed.filename == "file2"
+  assert parsed.version == 456
+
+
+@pytest.mark.parametrize(
+    "invalid_uri",
+    [
+        "http://example.com",
+        "artifact://invalid",
+        "artifact://app1/user1/sessions/session1/artifacts/file1",
+        "artifact://apps/app1/users/user1/sessions/session1/artifacts/file1",
+        "artifact://apps/app1/users/user1/artifacts/file1",
+    ],
+)
+def test_parse_invalid_artifact_uri(invalid_uri):
+  """Tests parsing invalid artifact URIs."""
+  assert artifact_util.parse_artifact_uri(invalid_uri) is None
+
+
+def test_get_session_scoped_artifact_uri():
+  """Tests constructing a session-scoped artifact URI."""
+  uri = artifact_util.get_artifact_uri(
+      app_name="app1",
+      user_id="user1",
+      session_id="session1",
+      filename="file1",
+      version=123,
+  )
+  assert (
+      uri
+      == "artifact://apps/app1/users/user1/sessions/session1/artifacts/file1/versions/123"
+  )
+
+
+def test_get_user_scoped_artifact_uri():
+  """Tests constructing a user-scoped artifact URI."""
+  uri = artifact_util.get_artifact_uri(
+      app_name="app2", user_id="user2", filename="file2", version=456
+  )
+  assert uri == "artifact://apps/app2/users/user2/artifacts/file2/versions/456"
+
+
+def test_is_artifact_ref_true():
+  """Tests is_artifact_ref with a valid artifact reference."""
+  artifact = types.Part(
+      file_data=types.FileData(
+          file_uri="artifact://apps/a/u/s/f/v/1", mime_type="text/plain"
+      )
+  )
+  assert artifact_util.is_artifact_ref(artifact) is True
+
+
+@pytest.mark.parametrize(
+    "part",
+    [
+        types.Part(text="hello"),
+        types.Part(inline_data=types.Blob(data=b"123", mime_type="text/plain")),
+        types.Part(
+            file_data=types.FileData(
+                file_uri="http://example.com", mime_type="text/plain"
+            )
+        ),
+        types.Part(),
+    ],
+)
+def test_is_artifact_ref_false(part):
+  """Tests is_artifact_ref with non-reference parts."""
+  assert artifact_util.is_artifact_ref(part) is False

tests/unittests/flows/llm_flows/test_contents.py

@@ -324,6 +324,63 @@ async def test_confirmation_events_are_filtered():
   ]
 
 
+@pytest.mark.asyncio
+async def test_rewind_events_are_filtered_out():
+  """Test that events are filtered based on rewind action."""
+  agent = Agent(model="gemini-2.5-flash", name="test_agent")
+  llm_request = LlmRequest(model="gemini-2.5-flash")
+  invocation_context = await testing_utils.create_invocation_context(
+      agent=agent
+  )
+
+  events = [
+      Event(
+          invocation_id="inv1",
+          author="user",
+          content=types.UserContent("First message"),
+      ),
+      Event(
+          invocation_id="inv1",
+          author="test_agent",
+          content=types.ModelContent("First response"),
+      ),
+      Event(
+          invocation_id="inv2",
+          author="user",
+          content=types.UserContent("Second message"),
+      ),
+      Event(
+          invocation_id="inv2",
+          author="test_agent",
+          content=types.ModelContent("Second response"),
+      ),
+      Event(
+          invocation_id="rewind_inv",
+          author="test_agent",
+          actions=EventActions(rewind_before_invocation_id="inv2"),
+      ),
+      Event(
+          invocation_id="inv3",
+          author="user",
+          content=types.UserContent("Third message"),
+      ),
+  ]
+  invocation_context.session.events = events
+
+  # Process the request
+  async for _ in contents.request_processor.run_async(
+      invocation_context, llm_request
+  ):
+    pass
+
+  # Verify rewind correctly filters conversation history
+  assert llm_request.contents == [
+      types.UserContent("First message"),
+      types.ModelContent("First response"),
+      types.UserContent("Third message"),
+  ]
+
+
 @pytest.mark.asyncio
 async def test_events_with_empty_content_are_skipped():
   """Test that events with empty content (state-only changes) are skipped."""

tests/unittests/runners/test_runner_rewind.py

@@ -0,0 +1,248 @@
+# Copyright 2025 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 google.adk.agents.base_agent import BaseAgent
+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 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 rewinded
+    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 rewinded 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 rewinded
+    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")