diff --git a/tests/unit/vertexai/genai/replays/test_create_agent_engine_session.py b/tests/unit/vertexai/genai/replays/test_create_agent_engine_session.py index b055a9b67c..3f3dc0e142 100644 --- a/tests/unit/vertexai/genai/replays/test_create_agent_engine_session.py +++ b/tests/unit/vertexai/genai/replays/test_create_agent_engine_session.py @@ -32,7 +32,7 @@ def test_create_session_with_ttl(client): config=types.CreateAgentEngineSessionConfig( display_name="my_session", session_state={"foo": "bar"}, - ttl="120s", + ttl="1200000s", labels={"label_key": "label_value"}, ), ) @@ -42,12 +42,13 @@ def test_create_session_with_ttl(client): assert operation.response.user_id == "test-user-123" assert operation.response.labels == {"label_key": "label_value"} assert operation.response.name.startswith(agent_engine.api_resource.name) + assert operation.done # Expire time is calculated by the server, so we only check that it is # within a reasonable range to avoid flakiness. assert ( - operation.response.create_time + datetime.timedelta(seconds=119.5) + operation.response.create_time + datetime.timedelta(seconds=1199999.5) <= operation.response.expire_time - <= operation.response.create_time + datetime.timedelta(seconds=120.5) + <= operation.response.create_time + datetime.timedelta(seconds=1200000.5) ) finally: # Clean up resources. @@ -60,7 +61,7 @@ def test_create_session_with_expire_time(client): assert isinstance(agent_engine, types.AgentEngine) assert isinstance(agent_engine.api_resource, types.ReasoningEngine) expire_time = datetime.datetime( - 2026, 1, 1, 12, 30, 00, tzinfo=datetime.timezone.utc + 2028, 1, 1, 12, 30, 00, tzinfo=datetime.timezone.utc ) operation = client.agent_engines.sessions.create( @@ -78,6 +79,7 @@ def test_create_session_with_expire_time(client): assert operation.response.user_id == "test-user-123" assert operation.response.name.startswith(agent_engine.api_resource.name) assert operation.response.expire_time == expire_time + assert operation.done finally: # Clean up resources. client.agent_engines.delete(name=agent_engine.api_resource.name, force=True) diff --git a/vertexai/_genai/sessions.py b/vertexai/_genai/sessions.py index 777bcf56a9..43c5bf64b5 100644 --- a/vertexai/_genai/sessions.py +++ b/vertexai/_genai/sessions.py @@ -651,12 +651,15 @@ def create( user_id=user_id, config=config, ) - if config.wait_for_completion and not operation.done: - operation = _agent_engines_utils._await_operation( - operation_name=operation.name, - get_operation_fn=self._get_session_operation, - poll_interval_seconds=0.5, - ) + if config.wait_for_completion: + if not operation.done: + operation = _agent_engines_utils._await_operation( + operation_name=operation.name, + get_operation_fn=self._get_session_operation, + poll_interval_seconds=0.5, + ) + # We need to make a call to get the session because the operation + # response might not contain the relevant fields. if operation.response: operation.response = self.get(name=operation.response.name) elif operation.error: @@ -1133,12 +1136,15 @@ async def create( user_id=user_id, config=config, ) - if config.wait_for_completion and not operation.done: - operation = await _agent_engines_utils._await_async_operation( - operation_name=operation.name, - get_operation_fn=self._get_session_operation, - poll_interval_seconds=0.5, - ) + if config.wait_for_completion: + if not operation.done: + operation = await _agent_engines_utils._await_async_operation( + operation_name=operation.name, + get_operation_fn=self._get_session_operation, + poll_interval_seconds=0.5, + ) + # We need to make a call to get the session because the operation + # response might not contain the relevant fields. if operation.response: operation.response = await self.get(name=operation.response.name) elif operation.error: