diff --git a/agentplatform/_genai/sessions.py b/agentplatform/_genai/sessions.py
index bcd4862494..97fd3a143f 100644
--- a/agentplatform/_genai/sessions.py
+++ b/agentplatform/_genai/sessions.py
@@ -625,7 +625,7 @@ def _update(
         *,
         name: str,
         config: Optional[types.UpdateAgentEngineSessionConfigOrDict] = None,
-    ) -> types.AgentEngineSessionOperation:
+    ) -> types.Session:
         """
         Updates an Agent Engine session.
 
@@ -636,7 +636,7 @@ def _update(
                 Optional. Additional configurations for updating the Agent Engine session.
 
         Returns:
-            AgentEngineSessionOperation: The operation for updating the Agent Engine session.
+            Session: The updated Agent Engine session.
 
         """
 
@@ -680,7 +680,7 @@ def _update(
 
         response_dict = {} if not response.body else json.loads(response.body)
 
-        return_value = types.AgentEngineSessionOperation._from_response(
+        return_value = types.Session._from_response(
             response=response_dict,
             kwargs=(
                 {
@@ -1222,7 +1222,7 @@ async def _update(
         *,
         name: str,
         config: Optional[types.UpdateAgentEngineSessionConfigOrDict] = None,
-    ) -> types.AgentEngineSessionOperation:
+    ) -> types.Session:
         """
         Updates an Agent Engine session.
 
@@ -1233,7 +1233,7 @@ async def _update(
                 Optional. Additional configurations for updating the Agent Engine session.
 
         Returns:
-            AgentEngineSessionOperation: The operation for updating the Agent Engine session.
+            Session: The updated Agent Engine session.
 
         """
 
@@ -1279,7 +1279,7 @@ async def _update(
 
         response_dict = {} if not response.body else json.loads(response.body)
 
-        return_value = types.AgentEngineSessionOperation._from_response(
+        return_value = types.Session._from_response(
             response=response_dict,
             kwargs=(
                 {
diff --git a/tests/unit/agentplatform/genai/replays/test_ae_session_private_update.py b/tests/unit/agentplatform/genai/replays/test_ae_session_private_update.py
index 3f62dee2f0..a2bb020b47 100644
--- a/tests/unit/agentplatform/genai/replays/test_ae_session_private_update.py
+++ b/tests/unit/agentplatform/genai/replays/test_ae_session_private_update.py
@@ -19,14 +19,14 @@
 
 
 def test_private_update_session(client):
-    agent_engine_session_operation = client.agent_engines.sessions._update(
+    session = client.agent_engines.sessions._update(
         name="reasoningEngines/2886612747586371584/sessions/3080649749292908544",
         config=types.UpdateAgentEngineSessionConfig(
             display_name="test-agent-engine-session-updated",
             user_id="test-user-id",
         ),
     )
-    assert isinstance(agent_engine_session_operation, types.AgentEngineSessionOperation)
+    assert isinstance(session, types.Session)
 
 
 pytestmark = pytest_helper.setup(