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

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

@@ -0,0 +1,59 @@
+# Copyright 2026 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 tests.unit.vertexai.genai.replays import pytest_helper
+from vertexai._genai import types
+
+
+def test_append_simple_a2a_task_events(client):
+    # Use the autopush environment.
+    client._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._api_client._http_options.api_version = "internal"
+    task = client.agent_engines.a2a_tasks.create(
+        name=agent_engine.api_resource.name,
+        a2a_task_id="task123",
+        context_id="context123",
+    )
+    assert isinstance(task, types.A2aTask)
+
+    client.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)
+
+
+pytestmark = pytest_helper.setup(
+    file=__file__,
+    globals_for_file=globals(),
+)

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

@@ -0,0 +1,50 @@
+# Copyright 2026 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 tests.unit.vertexai.genai.replays import pytest_helper
+from vertexai._genai import types
+
+
+def test_create_simple_a2a_task(client):
+    # Use the autopush environment.
+    client._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._api_client._http_options.api_version = "internal"
+
+    task = client.agent_engines.a2a_tasks.create(
+        name=agent_engine.api_resource.name,
+        a2a_task_id="task123",
+        context_id="context123",
+    )
+
+    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)
+
+
+pytestmark = pytest_helper.setup(
+    file=__file__,
+    globals_for_file=globals(),
+)

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

@@ -0,0 +1,54 @@
+# Copyright 2026 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 tests.unit.vertexai.genai.replays import pytest_helper
+from vertexai._genai import types
+
+
+def test_get_a2a_task(client):
+    # Use the autopush environment.
+    client._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._api_client._http_options.api_version = "internal"
+
+    created_task = client.agent_engines.a2a_tasks.create(
+        name=agent_engine.api_resource.name,
+        a2a_task_id="task123",
+        context_id="context123",
+    )
+    assert isinstance(created_task, types.A2aTask)
+
+    task = client.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)
+
+
+pytestmark = pytest_helper.setup(
+    file=__file__,
+    globals_for_file=globals(),
+)

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

@@ -0,0 +1,91 @@
+# Copyright 2026 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 tests.unit.vertexai.genai.replays import pytest_helper
+from vertexai._genai import types
+
+
+def test_list_simple_a2a_task_events(client):
+    # Use the autopush environment.
+    client._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._api_client._http_options.api_version = "internal"
+    task = client.agent_engines.a2a_tasks.create(
+        name=agent_engine.api_resource.name,
+        a2a_task_id="task999",
+        context_id="context999",
+    )
+    assert isinstance(task, types.A2aTask)
+
+    events = list(
+        client.agent_engines.a2a_tasks.events.list(
+            name=task.name,
+        )
+    )
+    assert len(events) == 1
+    assert events[0].event_data.state_change.new_state == "SUBMITTED"
+
+    client.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(
+        client.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)
+
+
+pytestmark = pytest_helper.setup(
+    file=__file__,
+    globals_for_file=globals(),
+)

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

@@ -0,0 +1,68 @@
+# Copyright 2026 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 tests.unit.vertexai.genai.replays import pytest_helper
+from vertexai._genai import types
+
+
+def test_list_a2a_tasks(client):
+    # Use the autopush environment.
+    client._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._api_client._http_options.api_version = "internal"
+
+    assert not list(
+        client.agent_engines.a2a_tasks.list(
+            name=agent_engine.api_resource.name,
+        )
+    )
+    client.agent_engines.a2a_tasks.create(
+        name=agent_engine.api_resource.name,
+        a2a_task_id="task123",
+        context_id="context123",
+    )
+    client.agent_engines.a2a_tasks.create(
+        name=agent_engine.api_resource.name,
+        a2a_task_id="task456",
+        context_id="context456",
+    )
+    a2a_tasks_list = client.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)
+
+
+pytestmark = pytest_helper.setup(
+    file=__file__,
+    globals_for_file=globals(),
+)