feat: Expose public update() and update_async() methods on Sessions and AsyncSessions client modules.

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.
+            types.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=(
                 {
@@ -770,6 +770,33 @@ def create(
                 )
         return operation
 
+    def update(
+        self,
+        *,
+        name: str,
+        config: Optional[types.UpdateAgentEngineSessionConfigOrDict] = None,
+    ) -> types.Session:
+        """Updates an Agent Engine session.
+
+        Args:
+            name (str):
+                Required. The name of the Agent Engine session to be updated. Format:
+                    `projects/{project}/locations/{location}/reasoningEngines/{resource_id}/sessions/{session_id}`.
+            config (UpdateAgentEngineSessionConfig):
+                Optional. The configuration for the session to update.
+
+        Returns:
+            Session: The updated Agent Engine session.
+        """
+        if config is None:
+            config = types.UpdateAgentEngineSessionConfig()
+        elif isinstance(config, dict):
+            config = types.UpdateAgentEngineSessionConfig.model_validate(config)
+        return self._update(
+            name=name,
+            config=config,
+        )
+
     def list(
         self,
         *,
@@ -1222,7 +1249,7 @@ async def _update(
         *,
         name: str,
         config: Optional[types.UpdateAgentEngineSessionConfigOrDict] = None,
-    ) -> types.AgentEngineSessionOperation:
+    ) -> types.Session:
         """
         Updates an Agent Engine session.
 
@@ -1233,7 +1260,7 @@ async def _update(
                 Optional. Additional configurations for updating the Agent Engine session.
 
         Returns:
-            AgentEngineSessionOperation: The operation for updating the Agent Engine session.
+            types.Session: The updated Agent Engine session.
 
         """
 
@@ -1279,7 +1306,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=(
                 {
@@ -1369,6 +1396,33 @@ async def create(
                 )
         return operation
 
+    async def update(
+        self,
+        *,
+        name: str,
+        config: Optional[types.UpdateAgentEngineSessionConfigOrDict] = None,
+    ) -> types.Session:
+        """Updates an Agent Engine session.
+
+        Args:
+            name (str):
+                Required. The name of the Agent Engine session to be updated. Format:
+                    `projects/{project}/locations/{location}/reasoningEngines/{resource_id}/sessions/{session_id}`.
+            config (UpdateAgentEngineSessionConfig):
+                Optional. The configuration for the session to update.
+
+        Returns:
+            Session: The updated Agent Engine session.
+        """
+        if config is None:
+            config = types.UpdateAgentEngineSessionConfig()
+        elif isinstance(config, dict):
+            config = types.UpdateAgentEngineSessionConfig.model_validate(config)
+        return await self._update(
+            name=name,
+            config=config,
+        )
+
     async def list(
         self,
         *,

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(

tests/unit/agentplatform/genai/replays/test_update_agent_engine_session.py

@@ -0,0 +1,114 @@
+# 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.
+#
+# pylint: disable=protected-access,bad-continuation,missing-function-docstring
+
+from agentplatform._genai import types
+from tests.unit.agentplatform.genai.replays import pytest_helper
+import pytest
+
+
+def test_update_session(client):
+  agent_engine = client.agent_engines.create()
+  try:
+    assert isinstance(agent_engine, types.AgentEngine)
+    assert isinstance(agent_engine.api_resource, types.ReasoningEngine)
+
+    session_operation = client.agent_engines.sessions.create(
+        name=agent_engine.api_resource.name,
+        user_id="test-user-123",
+        config=types.CreateAgentEngineSessionConfig(
+            display_name="initial_session",
+        ),
+    )
+    assert isinstance(session_operation, types.AgentEngineSessionOperation)
+
+    updated_session = client.agent_engines.sessions.update(
+        name=session_operation.response.name,
+        config=types.UpdateAgentEngineSessionConfig(
+            display_name="updated_session",
+            user_id="test-user-123",
+            labels={"env": "test", "tier": "dev"},
+        ),
+    )
+    assert isinstance(updated_session, types.Session)
+    assert updated_session.display_name == "updated_session"
+    assert updated_session.user_id == "test-user-123"
+    assert updated_session.labels == {"env": "test", "tier": "dev"}
+
+    # Second update: update with explicit update_mask
+    mask_updated_session = client.agent_engines.sessions.update(
+        name=session_operation.response.name,
+        config=types.UpdateAgentEngineSessionConfig(
+            display_name="session_with_mask",
+            user_id="test-user-123",
+            update_mask="displayName",
+        ),
+    )
+    assert isinstance(mask_updated_session, types.Session)
+    assert mask_updated_session.display_name == "session_with_mask"
+
+    # Third update: update with ttl (duration)
+    ttl_updated_session = client.agent_engines.sessions.update(
+        name=session_operation.response.name,
+        config=types.UpdateAgentEngineSessionConfig(
+            user_id="test-user-123",
+            ttl="86400s",
+        ),
+    )
+    assert isinstance(ttl_updated_session, types.Session)
+  finally:
+    client.agent_engines.delete(name=agent_engine.api_resource.name, force=True)
+
+
+pytestmark = pytest_helper.setup(
+    file=__file__,
+    globals_for_file=globals(),
+    test_method="agent_engines.sessions.update",
+)
+
+pytest_plugins = ("pytest_asyncio",)
+
+
+@pytest.mark.asyncio
+async def test_update_session_async(client):
+  agent_engine = client.agent_engines.create()
+  try:
+    assert isinstance(agent_engine, types.AgentEngine)
+    assert isinstance(agent_engine.api_resource, types.ReasoningEngine)
+
+    session_operation = await client.aio.agent_engines.sessions.create(
+        name=agent_engine.api_resource.name,
+        user_id="test-user-123",
+        config=types.CreateAgentEngineSessionConfig(
+            display_name="initial_session",
+        ),
+    )
+    assert isinstance(session_operation, types.AgentEngineSessionOperation)
+
+    updated_session = await client.aio.agent_engines.sessions.update(
+        name=session_operation.response.name,
+        config=types.UpdateAgentEngineSessionConfig(
+            display_name="updated_session",
+            user_id="test-user-123",
+            labels={"env": "test", "tier": "dev"},
+            ttl="86400s",
+        ),
+    )
+    assert isinstance(updated_session, types.Session)
+    assert updated_session.display_name == "updated_session"
+    assert updated_session.user_id == "test-user-123"
+    assert updated_session.labels == {"env": "test", "tier": "dev"}
+  finally:
+    client.agent_engines.delete(name=agent_engine.api_resource.name, force=True)