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

PiperOrigin-RevId: 820552460

Author: DeanChensj

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 @@ def _save_artifact(
       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 @@
 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 @@ async def save_artifact(
     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 @@ async def load_artifact(
       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):