fix: GenAI SDK client - Add get session call to create session sdk if an immediate success is returned

PiperOrigin-RevId: 881525769

Author: vertex-sdk-bot

tests/unit/vertexai/genai/replays/test_create_agent_engine_session.py

@@ -42,6 +42,7 @@ 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 (
@@ -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)

vertexai/_genai/sessions.py

@@ -651,13 +651,21 @@ 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,
-            )
+        logger.info("Create session operation response: %s", operation)
+        if config.wait_for_completion:
+            logger.info("Wait for completion")
+            if not operation.done:
+                logger.info("Poll operation")
+                operation = _agent_engines_utils._await_operation(
+                    operation_name=operation.name,
+                    get_operation_fn=self._get_session_operation,
+                    poll_interval_seconds=0.5,
+                )
+                logger.info("Poll operation done: %s", operation)
+            # We need to make a call to get the session because the operation
+            # response might not contain the relevant fields.
             if operation.response:
+                logger.info("Get session from operation response")
                 operation.response = self.get(name=operation.response.name)
             elif operation.error:
                 raise RuntimeError(f"Failed to create session: {operation.error}")
@@ -666,6 +674,7 @@ def create(
                     "Error retrieving session from the operation response. "
                     f"Operation name: {operation.name}"
                 )
+        logger.info("Create session final operation: %s", operation)
         return operation
 
     def list(