feat: add async wrappers in the GenAI client for the Agent Engine Task Store Service

vertex-sdk-bot · copybara-github · commit 199a40662a60 · 2026-03-03T16:03:29.000-08:00
PiperOrigin-RevId: 878165510
diff --git a/tests/unit/vertexai/genai/replays/test_append_agent_engine_a2a_task_events.py b/tests/unit/vertexai/genai/replays/test_append_agent_engine_a2a_task_events.py
@@ -16,6 +16,7 @@
 
 from tests.unit.vertexai.genai.replays import pytest_helper
 from vertexai._genai import types
+import pytest
 
 
 def test_append_simple_a2a_task_events(client):
@@ -57,3 +58,41 @@ def test_append_simple_a2a_task_events(client):
     file=__file__,
     globals_for_file=globals(),
 )
+
+pytest_plugins = ("pytest_asyncio",)
+
+
+@pytest.mark.asyncio
+async def test_append_simple_a2a_task_events_async(client):
+    # Use the autopush environment.
+    client.aio._api_client._http_options.base_url = (
+        "https://us-central1-autopush-aiplatform.sandbox.googleapis.com/"
+    )
+    agent_engine = client.agent_engines.create()
+    assert isinstance(agent_engine, types.AgentEngine)
+    assert isinstance(agent_engine.api_resource, types.ReasoningEngine)
+    # Use the internal API version for internal API access.
+    client.aio._api_client._http_options.api_version = "internal"
+    task = await client.aio.agent_engines.a2a_tasks.create(
+        name=agent_engine.api_resource.name,
+        a2a_task_id="task123",
+        config=types.CreateAgentEngineTaskConfig(context_id="context123"),
+    )
+    assert isinstance(task, types.A2aTask)
+
+    await client.aio.agent_engines.a2a_tasks.events.append(
+        name=task.name,
+        task_events=[
+            types.TaskEvent(
+                event_data=types.TaskEventData(
+                    metadata_change=types.TaskMetadataChange(
+                        new_metadata={"key1": "value1"}
+                    )
+                ),
+                event_sequence_number=1,
+            )
+        ],
+    )
+
+    # Clean up resources.
+    client.agent_engines.delete(name=agent_engine.api_resource.name, force=True)
diff --git a/tests/unit/vertexai/genai/replays/test_create_agent_engine_a2a_task.py b/tests/unit/vertexai/genai/replays/test_create_agent_engine_a2a_task.py
@@ -16,6 +16,7 @@
 
 from tests.unit.vertexai.genai.replays import pytest_helper
 from vertexai._genai import types
+import pytest
 
 
 def test_create_simple_a2a_task(client):
@@ -94,3 +95,78 @@ def test_create_simple_a2a_task(client):
     file=__file__,
     globals_for_file=globals(),
 )
+
+pytest_plugins = ("pytest_asyncio",)
+
+
+@pytest.mark.asyncio
+async def test_create_simple_a2a_task_async(client):
+    # Use the autopush environment.
+    client.aio._api_client._http_options.base_url = (
+        "https://us-central1-autopush-aiplatform.sandbox.googleapis.com/"
+    )
+    agent_engine = client.agent_engines.create()
+    assert isinstance(agent_engine, types.AgentEngine)
+    assert isinstance(agent_engine.api_resource, types.ReasoningEngine)
+    # Use the internal API version for internal API access.
+    client.aio._api_client._http_options.api_version = "internal"
+
+    task = await client.aio.agent_engines.a2a_tasks.create(
+        name=agent_engine.api_resource.name,
+        a2a_task_id="task123",
+        config=types.CreateAgentEngineTaskConfig(
+            context_id="context123",
+            metadata={
+                "key": "value",
+                "key2": [{"key3": "value3", "key4": "value4"}],
+            },
+            status_details=types.TaskStatusDetails(
+                task_message=types.TaskMessage(
+                    role="user",
+                    message_id="message123",
+                    parts=[
+                        types.Part(
+                            text="hello123",
+                        )
+                    ],
+                    metadata={
+                        "key42": "value42",
+                    },
+                ),
+            ),
+            output=types.TaskOutput(
+                artifacts=[
+                    types.TaskArtifact(
+                        artifact_id="artifact123",
+                        display_name="display_name123",
+                        description="description123",
+                        parts=[
+                            types.Part(
+                                text="hello456",
+                            )
+                        ],
+                    )
+                ],
+            ),
+        ),
+    )
+
+    assert isinstance(task, types.A2aTask)
+    assert task.name == f"{agent_engine.api_resource.name}/a2aTasks/task123"
+    assert task.context_id == "context123"
+    assert task.state == types.State.SUBMITTED
+    assert task.status_details.task_message.role == "user"
+    assert task.status_details.task_message.message_id == "message123"
+    assert task.status_details.task_message.parts[0].text == "hello123"
+    assert task.status_details.task_message.metadata["key42"] == "value42"
+    assert task.output.artifacts[0].artifact_id == "artifact123"
+    assert task.output.artifacts[0].display_name == "display_name123"
+    assert task.output.artifacts[0].description == "description123"
+    assert task.output.artifacts[0].parts[0].text == "hello456"
+    assert task.metadata == {
+        "key": "value",
+        "key2": [{"key3": "value3", "key4": "value4"}],
+    }
+
+    # Clean up resources.
+    client.agent_engines.delete(name=agent_engine.api_resource.name, force=True)
diff --git a/tests/unit/vertexai/genai/replays/test_get_agent_engine_a2a_task.py b/tests/unit/vertexai/genai/replays/test_get_agent_engine_a2a_task.py
@@ -16,6 +16,7 @@
 
 from tests.unit.vertexai.genai.replays import pytest_helper
 from vertexai._genai import types
+import pytest
 
 
 def test_get_a2a_task(client):
@@ -52,3 +53,37 @@ def test_get_a2a_task(client):
     file=__file__,
     globals_for_file=globals(),
 )
+
+
+pytest_plugins = ("pytest_asyncio",)
+
+
+@pytest.mark.asyncio
+async def test_get_a2a_task_async(client):
+    # Use the autopush environment.
+    client.aio._api_client._http_options.base_url = (
+        "https://us-central1-autopush-aiplatform.sandbox.googleapis.com/"
+    )
+    agent_engine = client.agent_engines.create()
+    assert isinstance(agent_engine, types.AgentEngine)
+    assert isinstance(agent_engine.api_resource, types.ReasoningEngine)
+    # Use the internal API version for internal API access.
+    client.aio._api_client._http_options.api_version = "internal"
+
+    created_task = await client.aio.agent_engines.a2a_tasks.create(
+        name=agent_engine.api_resource.name,
+        a2a_task_id="task123",
+        config=types.CreateAgentEngineTaskConfig(context_id="context123"),
+    )
+    assert isinstance(created_task, types.A2aTask)
+
+    task = await client.aio.agent_engines.a2a_tasks.get(
+        name=created_task.name,
+    )
+    assert isinstance(task, types.A2aTask)
+    assert task.name == f"{agent_engine.api_resource.name}/a2aTasks/task123"
+    assert task.context_id == "context123"
+    assert task.state == types.State.SUBMITTED
+
+    # Clean up resources.
+    client.agent_engines.delete(name=agent_engine.api_resource.name, force=True)
diff --git a/tests/unit/vertexai/genai/replays/test_list_agent_engine_a2a_task_events.py b/tests/unit/vertexai/genai/replays/test_list_agent_engine_a2a_task_events.py
@@ -16,6 +16,7 @@
 
 from tests.unit.vertexai.genai.replays import pytest_helper
 from vertexai._genai import types
+import pytest
 
 
 def test_list_simple_a2a_task_events(client):
@@ -89,3 +90,71 @@ def test_list_simple_a2a_task_events(client):
     file=__file__,
     globals_for_file=globals(),
 )
+
+
+@pytest.mark.asyncio
+async def test_list_simple_a2a_task_events_async(client):
+    # Use the autopush environment.
+    client.aio._api_client._http_options.base_url = (
+        "https://us-central1-autopush-aiplatform.sandbox.googleapis.com/"
+    )
+    agent_engine = client.agent_engines.create()
+    assert isinstance(agent_engine, types.AgentEngine)
+    assert isinstance(agent_engine.api_resource, types.ReasoningEngine)
+    # Use the internal API version for internal API access.
+    client.aio._api_client._http_options.api_version = "internal"
+    task = await client.aio.agent_engines.a2a_tasks.create(
+        name=agent_engine.api_resource.name,
+        a2a_task_id="task999",
+        config=types.CreateAgentEngineTaskConfig(context_id="context999"),
+    )
+    assert isinstance(task, types.A2aTask)
+
+    events = list(
+        await client.aio.agent_engines.a2a_tasks.events.list(
+            name=task.name,
+        )
+    )
+    assert len(events) == 1
+    assert events[0].event_data.state_change.new_state == "SUBMITTED"
+
+    await client.aio.agent_engines.a2a_tasks.events.append(
+        name=task.name,
+        task_events=[
+            types.TaskEvent(
+                event_data=types.TaskEventData(
+                    metadata_change=types.TaskMetadataChange(
+                        new_metadata={"key1": "value1"}
+                    )
+                ),
+                event_sequence_number=1,
+            ),
+            types.TaskEvent(
+                event_data=types.TaskEventData(
+                    metadata_change=types.TaskMetadataChange(
+                        new_metadata={"key2": "value2"}
+                    )
+                ),
+                event_sequence_number=2,
+            ),
+        ],
+    )
+
+    result = list(
+        await client.aio.agent_engines.a2a_tasks.events.list(
+            name=task.name,
+            config=types.ListAgentEngineTaskEventsConfig(
+                order_by="create_time desc",
+            ),
+        )
+    )
+
+    assert len(result) == 3
+    assert result[0].event_sequence_number == 2
+    assert result[0].event_data.metadata_change.new_metadata == {"key2": "value2"}
+    assert result[1].event_sequence_number == 1
+    assert result[1].event_data.metadata_change.new_metadata == {"key1": "value1"}
+    assert result[2].event_data.state_change.new_state == "SUBMITTED"
+
+    # Clean up resources.
+    client.agent_engines.delete(name=agent_engine.api_resource.name, force=True)
diff --git a/tests/unit/vertexai/genai/replays/test_list_agent_engine_a2a_tasks.py b/tests/unit/vertexai/genai/replays/test_list_agent_engine_a2a_tasks.py
@@ -16,6 +16,7 @@
 
 from tests.unit.vertexai.genai.replays import pytest_helper
 from vertexai._genai import types
+import pytest
 
 
 def test_list_a2a_tasks(client):
@@ -66,3 +67,51 @@ def test_list_a2a_tasks(client):
     file=__file__,
     globals_for_file=globals(),
 )
+
+
+pytest_plugins = ("pytest_asyncio",)
+
+
+@pytest.mark.asyncio
+async def test_list_a2a_tasks_async(client):
+    # Use the autopush environment.
+    client.aio._api_client._http_options.base_url = (
+        "https://us-central1-autopush-aiplatform.sandbox.googleapis.com/"
+    )
+    agent_engine = client.agent_engines.create()
+    assert isinstance(agent_engine, types.AgentEngine)
+    assert isinstance(agent_engine.api_resource, types.ReasoningEngine)
+    # Use the internal API version for internal API access.
+    client.aio._api_client._http_options.api_version = "internal"
+
+    assert not list(
+        await client.aio.agent_engines.a2a_tasks.list(
+            name=agent_engine.api_resource.name,
+        )
+    )
+    await client.aio.agent_engines.a2a_tasks.create(
+        name=agent_engine.api_resource.name,
+        a2a_task_id="task123",
+        config=types.CreateAgentEngineTaskConfig(context_id="context123"),
+    )
+    await client.aio.agent_engines.a2a_tasks.create(
+        name=agent_engine.api_resource.name,
+        a2a_task_id="task456",
+        config=types.CreateAgentEngineTaskConfig(context_id="context456"),
+    )
+    a2a_tasks_list = await client.aio.agent_engines.a2a_tasks.list(
+        name=agent_engine.api_resource.name,
+        config=types.ListAgentEngineTasksConfig(
+            page_size=1,
+            order_by="create_time asc",
+        ),
+    )
+    assert len(a2a_tasks_list) == 1
+    assert isinstance(a2a_tasks_list[0], types.A2aTask)
+    assert a2a_tasks_list[0].name == (
+        f"{agent_engine.api_resource.name}/a2aTasks/task123"
+    )
+    assert a2a_tasks_list[0].context_id == "context123"
+
+    # Clean up resources.
+    agent_engine.delete(force=True)
diff --git a/vertexai/_genai/a2a_task_events.py b/vertexai/_genai/a2a_task_events.py
@@ -25,7 +25,7 @@
 from google.genai import _common
 from google.genai._common import get_value_by_path as getv
 from google.genai._common import set_value_by_path as setv
-from google.genai.pagers import Pager
+from google.genai.pagers import AsyncPager, Pager
 
 from . import types
 
@@ -245,8 +245,6 @@ def _list(
         self._api_client._verify_response(return_value)
         return return_value
 
-    """Task events."""
-
     def list(
         self,
         *,
@@ -415,3 +413,28 @@ async def _list(
 
         self._api_client._verify_response(return_value)
         return return_value
+
+    async def list(
+        self,
+        *,
+        name: str,
+        config: Optional[types.ListAgentEngineTaskEventsConfigOrDict] = None,
+    ) -> AsyncPager[types.TaskEvent]:
+        """Lists the A2A tasks of an Agent Engine.
+
+        Args:
+            name (str):
+                Required. The name of the agent engine to list tasks for.
+            config (List):
+                Optional. The configuration for the tasks to list.
+
+        Returns:
+            AsyncPager[TaskEvent]: An async pager of Task events.
+        """
+
+        return AsyncPager(
+            "taskEvents",
+            functools.partial(self._list, name=name),
+            await self._list(name=name, config=config),
+            config,
+        )
diff --git a/vertexai/_genai/a2a_tasks.py b/vertexai/_genai/a2a_tasks.py
diff --git a/vertexai/_genai/agent_engines.py b/vertexai/_genai/agent_engines.py
