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

PiperOrigin-RevId: 874762764

Author: vertex-sdk-bot

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",
+        config=types.CreateAgentEngineTaskConfig(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,97 @@
+# 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",
+      config=types.CreateAgentEngineTaskConfig(
+          context_id="context123",
+          metadata={
+              "key": "value",
+              "key2": [{"key3": "value3", "key4": "value4"}],
+          },
+          state=types.A2aTaskState.STATE_UNSPECIFIED,
+          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)
+
+
+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",
+        config=types.CreateAgentEngineTaskConfig(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",
+        config=types.CreateAgentEngineTaskConfig(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",
+        config=types.CreateAgentEngineTaskConfig(context_id="context123"),
+    )
+    client.agent_engines.a2a_tasks.create(
+        name=agent_engine.api_resource.name,
+        a2a_task_id="task456",
+        config=types.CreateAgentEngineTaskConfig(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(),
+)