feat: Add custom ID generators to SimpleRequestContextBuilder (#594)

# Description

This change allows passing custom `task_id_generator` and
`context_id_generator` functions to the `SimpleRequestContextBuilder`.

This provides flexibility in how task and context IDs are generated,
defaulting to the previous behavior if no generators are provided.

Thank you for opening a Pull Request!
Before submitting your PR, there are a few things you can do to make
sure it goes smoothly:

- [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 #<issue_number_goes_here> 🦕

Author: chenweiyang0204

src/a2a/server/agent_execution/simple_request_context_builder.py

@@ -2,6 +2,7 @@
 
 from a2a.server.agent_execution import RequestContext, RequestContextBuilder
 from a2a.server.context import ServerCallContext
+from a2a.server.id_generator import IDGenerator
 from a2a.server.tasks import TaskStore
 from a2a.types import MessageSendParams, Task
 
@@ -13,6 +14,8 @@ def __init__(
         self,
         should_populate_referred_tasks: bool = False,
         task_store: TaskStore | None = None,
+        task_id_generator: IDGenerator | None = None,
+        context_id_generator: IDGenerator | None = None,
     ) -> None:
         """Initializes the SimpleRequestContextBuilder.
 
@@ -22,9 +25,13 @@ def __init__(
                 `related_tasks` field in the RequestContext. Defaults to False.
             task_store: The TaskStore instance to use for fetching referred tasks.
                 Required if `should_populate_referred_tasks` is True.
+            task_id_generator: ID generator for new task IDs. Defaults to None.
+            context_id_generator: ID generator for new context IDs. Defaults to None.
         """
         self._task_store = task_store
         self._should_populate_referred_tasks = should_populate_referred_tasks
+        self._task_id_generator = task_id_generator
+        self._context_id_generator = context_id_generator
 
     async def build(
         self,
@@ -74,4 +81,6 @@ async def build(
             task=task,
             related_tasks=related_tasks,
             call_context=context,
+            task_id_generator=self._task_id_generator,
+            context_id_generator=self._context_id_generator,
         )

tests/server/agent_execution/test_simple_request_context_builder.py

@@ -10,6 +10,7 @@
     SimpleRequestContextBuilder,
 )
 from a2a.server.context import ServerCallContext
+from a2a.server.id_generator import IDGenerator
 from a2a.server.tasks.task_store import TaskStore
 from a2a.types import (
     Message,
@@ -275,6 +276,65 @@ async def test_build_populate_false_with_reference_task_ids(self) -> None:
         self.assertEqual(request_context.related_tasks, [])
         self.mock_task_store.get.assert_not_called()
 
+    async def test_build_with_custom_id_generators(self) -> None:
+        mock_task_id_generator = AsyncMock(spec=IDGenerator)
+        mock_context_id_generator = AsyncMock(spec=IDGenerator)
+        mock_task_id_generator.generate.return_value = 'custom_task_id'
+        mock_context_id_generator.generate.return_value = 'custom_context_id'
+
+        builder = SimpleRequestContextBuilder(
+            should_populate_referred_tasks=False,
+            task_store=self.mock_task_store,
+            task_id_generator=mock_task_id_generator,
+            context_id_generator=mock_context_id_generator,
+        )
+        params = MessageSendParams(message=create_sample_message())
+        server_call_context = ServerCallContext(user=UnauthenticatedUser())
+
+        request_context = await builder.build(
+            params=params,
+            task_id=None,
+            context_id=None,
+            task=None,
+            context=server_call_context,
+        )
+
+        mock_task_id_generator.generate.assert_called_once()
+        mock_context_id_generator.generate.assert_called_once()
+        self.assertEqual(request_context.task_id, 'custom_task_id')
+        self.assertEqual(request_context.context_id, 'custom_context_id')
+
+    async def test_build_with_provided_ids_and_custom_id_generators(
+        self,
+    ) -> None:
+        mock_task_id_generator = AsyncMock(spec=IDGenerator)
+        mock_context_id_generator = AsyncMock(spec=IDGenerator)
+
+        builder = SimpleRequestContextBuilder(
+            should_populate_referred_tasks=False,
+            task_store=self.mock_task_store,
+            task_id_generator=mock_task_id_generator,
+            context_id_generator=mock_context_id_generator,
+        )
+        params = MessageSendParams(message=create_sample_message())
+        server_call_context = ServerCallContext(user=UnauthenticatedUser())
+
+        provided_task_id = 'provided_task_id'
+        provided_context_id = 'provided_context_id'
+
+        request_context = await builder.build(
+            params=params,
+            task_id=provided_task_id,
+            context_id=provided_context_id,
+            task=None,
+            context=server_call_context,
+        )
+
+        mock_task_id_generator.generate.assert_not_called()
+        mock_context_id_generator.generate.assert_not_called()
+        self.assertEqual(request_context.task_id, provided_task_id)
+        self.assertEqual(request_context.context_id, provided_context_id)
+
 
 if __name__ == '__main__':
     unittest.main()