feat(api_server): expose session rewind as a REST endpoint

src/google/adk/cli/api_server.py

@@ -622,6 +622,13 @@ class UpdateSessionRequest(common.BaseModel):
   """The state changes to apply to the session."""
 
 
+class RewindSessionRequest(common.BaseModel):
+  """Request to rewind a session to before a given invocation."""
+
+  rewind_before_invocation_id: str
+  """The invocation ID to rewind the session to before."""
+
+
 class FinalizeAgentIdentityCredentialsRequest(common.BaseModel):
   """Request to finalize a 3LO consent for an Agent Identity connector."""
 
@@ -1590,6 +1597,47 @@ async def update_session(
 
       return session
 
+    @app.post(
+        "/apps/{app_name}/users/{user_id}/sessions/{session_id}/rewind",
+        response_model_exclude_none=True,
+    )
+    async def rewind_session(
+        app_name: str,
+        user_id: str,
+        session_id: str,
+        req: RewindSessionRequest,
+    ) -> Session:
+      """Rewinds a session to before the specified invocation.
+
+      Args:
+          app_name: The name of the application.
+          user_id: The ID of the user.
+          session_id: The ID of the session to rewind.
+          req: The rewind request identifying the invocation to rewind to.
+
+      Returns:
+          The rewound session.
+
+      Raises:
+          HTTPException: If the session or invocation is not found.
+      """
+      runner = await self.get_runner_async(app_name)
+      try:
+        await runner.rewind_async(
+            user_id=user_id,
+            session_id=session_id,
+            rewind_before_invocation_id=req.rewind_before_invocation_id,
+        )
+      except ValueError as ve:
+        raise HTTPException(status_code=404, detail=str(ve)) from ve
+
+      session = await self.session_service.get_session(
+          app_name=app_name, user_id=user_id, session_id=session_id
+      )
+      if not session:
+        raise HTTPException(status_code=404, detail="Session not found")
+      return session
+
     @app.get(
         "/apps/{app_name}/users/{user_id}/sessions/{session_id}/artifacts/{artifact_name:path}/versions/{version_id}/metadata",
         response_model=ArtifactVersion,

tests/unittests/cli/test_fast_api.py

@@ -1737,6 +1737,60 @@ def test_patch_session_not_found(test_app, test_session_info):
   logger.info("Patch session not found test passed")
 
 
+def test_rewind_session(test_app, create_test_session, monkeypatch):
+  """Test rewinding a session to before a given invocation."""
+  info = create_test_session
+  captured: dict[str, Any] = {}
+
+  async def rewind_async_capture(
+      self,
+      *,
+      user_id: str,
+      session_id: str,
+      rewind_before_invocation_id: str,
+      run_config: Optional[RunConfig] = None,
+  ):
+    del self, run_config
+    captured["user_id"] = user_id
+    captured["session_id"] = session_id
+    captured["rewind_before_invocation_id"] = rewind_before_invocation_id
+
+  monkeypatch.setattr(Runner, "rewind_async", rewind_async_capture)
+
+  url = f"/apps/{info['app_name']}/users/{info['user_id']}/sessions/{info['session_id']}/rewind"
+  response = test_app.post(
+      url, json={"rewind_before_invocation_id": "some-invocation-id"}
+  )
+
+  assert response.status_code == 200
+  assert captured["user_id"] == info["user_id"]
+  assert captured["session_id"] == info["session_id"]
+  assert captured["rewind_before_invocation_id"] == "some-invocation-id"
+  data = response.json()
+  assert data["id"] == info["session_id"]
+
+
+def test_rewind_session_invocation_not_found(
+    test_app, create_test_session, monkeypatch
+):
+  """Test rewinding to an unknown invocation ID returns 404."""
+  info = create_test_session
+
+  async def rewind_async_raise(self, **kwargs):
+    del self, kwargs
+    raise ValueError("Invocation ID not found: missing-invocation-id")
+
+  monkeypatch.setattr(Runner, "rewind_async", rewind_async_raise)
+
+  url = f"/apps/{info['app_name']}/users/{info['user_id']}/sessions/{info['session_id']}/rewind"
+  response = test_app.post(
+      url, json={"rewind_before_invocation_id": "missing-invocation-id"}
+  )
+
+  assert response.status_code == 404
+  assert "missing-invocation-id" in response.json()["detail"]
+
+
 def test_agent_run(test_app, create_test_session):
   """Test running an agent with a message."""
   info = create_test_session