feat: custom ID generators (#490)

# Description

This PR introduces a new feature allowing server maintainers to provide
custom ID generators for system-created objects.

I identified 2 key areas requiring ID generation:
-
[`ResourceContext`](https://github.com/a2aproject/a2a-python/blob/19091262d07b5226becdb177ed4d780b4befb5f4/src/a2a/server/agent_execution/context.py):
Manages the creation of new task and context IDs when not supplied by
the client.
-
[`TaskUpdater`](https://github.com/a2aproject/a2a-python/blob/19091262d07b5226becdb177ed4d780b4befb5f4/src/a2a/server/tasks/task_updater.py):
Helper responsible for generating IDs for new artifacts and messages
during task updates.

To address this, a new abstract class, `IDGenerator`, has been
introduced. This class enables the creation of specific ID generators
for different resource types, which can then be integrated into the
relevant SDK components.

Alternative approaches were considered:
1. Using a single `IDGenerator` instance for all resource types, with
resource type information passed to the generate method. This option was
discarded due to potential coupling issues and increased maintenance
complexity.
1. Implementing an additional `IDGeneratorRegistry` to manage
per-resource-type generators via a `get_id_generator(resource_type)`
method. While this would simplify object passing, it introduces
additional indirection and overall system complexity.

# Prerequisites
- [x] Follow the [`CONTRIBUTING`
Guide](https://github.com/a2aproject/a2a-python/blob/main/CONTRIBUTING.md).
- [x] Make your Pull Request title in the
<https://www.conventionalcommits.org/> specification.
- Important Prefixes for
[release-please](https://github.com/googleapis/release-please):
- `fix:` which represents bug fixes, and correlates to a
[SemVer](https://semver.org/) patch.
- `feat:` represents a new feature, and correlates to a SemVer minor.
- `feat!:`, or `fix!:`, `refactor!:`, etc., which represent a breaking
change (indicated by the `!`) and will result in a SemVer major.
- [x] Ensure the tests and linter pass (Run `bash scripts/format.sh`
from the repository root to format)
- [x] Appropriate docs were updated (if necessary)

Fixes #378 🦕

---------

Co-authored-by: lkawka <[email protected]>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Co-authored-by: Holt Skinner <[email protected]>

Author: 4 people

src/a2a/server/agent_execution/context.py

@@ -1,8 +1,11 @@
-import uuid
-
 from typing import Any
 
 from a2a.server.context import ServerCallContext
+from a2a.server.id_generator import (
+    IDGenerator,
+    IDGeneratorContext,
+    UUIDGenerator,
+)
 from a2a.types import (
     InvalidParamsError,
     Message,
@@ -30,6 +33,8 @@ def __init__(  # noqa: PLR0913
         task: Task | None = None,
         related_tasks: list[Task] | None = None,
         call_context: ServerCallContext | None = None,
+        task_id_generator: IDGenerator | None = None,
+        context_id_generator: IDGenerator | None = None,
     ):
         """Initializes the RequestContext.
 
@@ -40,6 +45,8 @@ def __init__(  # noqa: PLR0913
             task: The existing `Task` object retrieved from the store, if any.
             related_tasks: A list of other tasks related to the current request (e.g., for tool use).
             call_context: The server call context associated with this request.
+            task_id_generator: ID generator for new task IDs. Defaults to UUID generator.
+            context_id_generator: ID generator for new context IDs. Defaults to UUID generator.
         """
         if related_tasks is None:
             related_tasks = []
@@ -49,6 +56,12 @@ def __init__(  # noqa: PLR0913
         self._current_task = task
         self._related_tasks = related_tasks
         self._call_context = call_context
+        self._task_id_generator = (
+            task_id_generator if task_id_generator else UUIDGenerator()
+        )
+        self._context_id_generator = (
+            context_id_generator if context_id_generator else UUIDGenerator()
+        )
         # If the task id and context id were provided, make sure they
         # match the request. Otherwise, create them
         if self._params:
@@ -163,7 +176,9 @@ def _check_or_generate_task_id(self) -> None:
             return
 
         if not self._task_id and not self._params.message.task_id:
-            self._params.message.task_id = str(uuid.uuid4())
+            self._params.message.task_id = self._task_id_generator.generate(
+                IDGeneratorContext(context_id=self._context_id)
+            )
         if self._params.message.task_id:
             self._task_id = self._params.message.task_id
 
@@ -173,6 +188,10 @@ def _check_or_generate_context_id(self) -> None:
             return
 
         if not self._context_id and not self._params.message.context_id:
-            self._params.message.context_id = str(uuid.uuid4())
+            self._params.message.context_id = (
+                self._context_id_generator.generate(
+                    IDGeneratorContext(task_id=self._task_id)
+                )
+            )
         if self._params.message.context_id:
             self._context_id = self._params.message.context_id

src/a2a/server/id_generator.py

@@ -0,0 +1,28 @@
+import uuid
+
+from abc import ABC, abstractmethod
+
+from pydantic import BaseModel
+
+
+class IDGeneratorContext(BaseModel):
+    """Context for providing additional information to ID generators."""
+
+    task_id: str | None = None
+    context_id: str | None = None
+
+
+class IDGenerator(ABC):
+    """Interface for generating unique identifiers."""
+
+    @abstractmethod
+    def generate(self, context: IDGeneratorContext) -> str:
+        pass
+
+
+class UUIDGenerator(IDGenerator):
+    """UUID implementation of the IDGenerator interface."""
+
+    def generate(self, context: IDGeneratorContext) -> str:
+        """Generates a random UUID, ignoring the context."""
+        return str(uuid.uuid4())

src/a2a/server/tasks/task_updater.py

@@ -1,10 +1,14 @@
 import asyncio
-import uuid
 
 from datetime import datetime, timezone
 from typing import Any
 
 from a2a.server.events import EventQueue
+from a2a.server.id_generator import (
+    IDGenerator,
+    IDGeneratorContext,
+    UUIDGenerator,
+)
 from a2a.types import (
     Artifact,
     Message,
@@ -23,13 +27,22 @@ class TaskUpdater:
     Simplifies the process of creating and enqueueing standard task events.
     """
 
-    def __init__(self, event_queue: EventQueue, task_id: str, context_id: str):
+    def __init__(
+        self,
+        event_queue: EventQueue,
+        task_id: str,
+        context_id: str,
+        artifact_id_generator: IDGenerator | None = None,
+        message_id_generator: IDGenerator | None = None,
+    ):
         """Initializes the TaskUpdater.
 
         Args:
             event_queue: The `EventQueue` associated with the task.
             task_id: The ID of the task.
             context_id: The context ID of the task.
+            artifact_id_generator: ID generator for new artifact IDs. Defaults to UUID generator.
+            message_id_generator: ID generator for new message IDs. Defaults to UUID generator.
         """
         self.event_queue = event_queue
         self.task_id = task_id
@@ -42,6 +55,12 @@ def __init__(self, event_queue: EventQueue, task_id: str, context_id: str):
             TaskState.failed,
             TaskState.rejected,
         }
+        self._artifact_id_generator = (
+            artifact_id_generator if artifact_id_generator else UUIDGenerator()
+        )
+        self._message_id_generator = (
+            message_id_generator if message_id_generator else UUIDGenerator()
+        )
 
     async def update_status(
         self,
@@ -110,7 +129,11 @@ async def add_artifact(  # noqa: PLR0913
             extensions: Optional list of extensions for the artifact.
         """
         if not artifact_id:
-            artifact_id = str(uuid.uuid4())
+            artifact_id = self._artifact_id_generator.generate(
+                IDGeneratorContext(
+                    task_id=self.task_id, context_id=self.context_id
+                )
+            )
 
         await self.event_queue.enqueue_event(
             TaskArtifactUpdateEvent(
@@ -205,7 +228,11 @@ def new_agent_message(
             role=Role.agent,
             task_id=self.task_id,
             context_id=self.context_id,
-            message_id=str(uuid.uuid4()),
+            message_id=self._message_id_generator.generate(
+                IDGeneratorContext(
+                    task_id=self.task_id, context_id=self.context_id
+                )
+            ),
             metadata=metadata,
             parts=parts,
         )

tests/server/agent_execution/test_context.py

@@ -6,6 +6,7 @@
 
 from a2a.server.agent_execution import RequestContext
 from a2a.server.context import ServerCallContext
+from a2a.server.id_generator import IDGenerator
 from a2a.types import (
     Message,
     MessageSendParams,
@@ -149,6 +150,20 @@ def test_check_or_generate_task_id_with_existing_task_id(self, mock_params):
         assert context.task_id == existing_id
         assert mock_params.message.task_id == existing_id
 
+    def test_check_or_generate_task_id_with_custom_id_generator(
+        self, mock_params
+    ):
+        """Test _check_or_generate_task_id uses custom ID generator when provided."""
+        id_generator = Mock(spec=IDGenerator)
+        id_generator.generate.return_value = 'custom-task-id'
+
+        context = RequestContext(
+            request=mock_params, task_id_generator=id_generator
+        )
+        # The method is called during initialization
+
+        assert context.task_id == 'custom-task-id'
+
     def test_check_or_generate_context_id_no_params(self):
         """Test _check_or_generate_context_id with no params does nothing."""
         context = RequestContext()
@@ -168,6 +183,20 @@ def test_check_or_generate_context_id_with_existing_context_id(
         assert context.context_id == existing_id
         assert mock_params.message.context_id == existing_id
 
+    def test_check_or_generate_context_id_with_custom_id_generator(
+        self, mock_params
+    ):
+        """Test _check_or_generate_context_id uses custom ID generator when provided."""
+        id_generator = Mock(spec=IDGenerator)
+        id_generator.generate.return_value = 'custom-context-id'
+
+        context = RequestContext(
+            request=mock_params, context_id_generator=id_generator
+        )
+        # The method is called during initialization
+
+        assert context.context_id == 'custom-context-id'
+
     def test_init_raises_error_on_task_id_mismatch(
         self, mock_params, mock_task
     ):

tests/server/tasks/test_task_updater.py

@@ -1,11 +1,12 @@
 import asyncio
 import uuid
 
-from unittest.mock import AsyncMock, patch
+from unittest.mock import AsyncMock, Mock, patch
 
 import pytest
 
 from a2a.server.events import EventQueue
+from a2a.server.id_generator import IDGenerator
 from a2a.server.tasks import TaskUpdater
 from a2a.types import (
     Message,
@@ -151,6 +152,26 @@ async def test_add_artifact_generates_id(
     assert event.last_chunk is None
 
 
+@pytest.mark.asyncio
+async def test_add_artifact_generates_custom_id(event_queue, sample_parts):
+    """Test add_artifact uses a custom ID generator when provided."""
+    artifact_id_generator = Mock(spec=IDGenerator)
+    artifact_id_generator.generate.return_value = 'custom-artifact-id'
+    task_updater = TaskUpdater(
+        event_queue=event_queue,
+        task_id='test-task-id',
+        context_id='test-context-id',
+        artifact_id_generator=artifact_id_generator,
+    )
+
+    await task_updater.add_artifact(parts=sample_parts, artifact_id=None)
+
+    event_queue.enqueue_event.assert_called_once()
+    event = event_queue.enqueue_event.call_args[0][0]
+    assert isinstance(event, TaskArtifactUpdateEvent)
+    assert event.artifact.artifact_id == 'custom-artifact-id'
+
+
 @pytest.mark.asyncio
 @pytest.mark.parametrize(
     'append_val, last_chunk_val',
@@ -304,6 +325,22 @@ def test_new_agent_message_with_metadata(task_updater, sample_parts):
     assert message.metadata == metadata
 
 
+def test_new_agent_message_with_custom_id_generator(event_queue, sample_parts):
+    """Test creating a new agent message with a custom message ID generator."""
+    message_id_generator = Mock(spec=IDGenerator)
+    message_id_generator.generate.return_value = 'custom-message-id'
+    task_updater = TaskUpdater(
+        event_queue=event_queue,
+        task_id='test-task-id',
+        context_id='test-context-id',
+        message_id_generator=message_id_generator,
+    )
+
+    message = task_updater.new_agent_message(parts=sample_parts)
+
+    assert message.message_id == 'custom-message-id'
+
+
 @pytest.mark.asyncio
 async def test_failed_without_message(task_updater, event_queue):
     """Test marking a task as failed without a message."""