diff --git a/src/a2a/server/agent_execution/simple_request_context_builder.py b/src/a2a/server/agent_execution/simple_request_context_builder.py index 3eca4435..876b6561 100644 --- a/src/a2a/server/agent_execution/simple_request_context_builder.py +++ b/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, ) diff --git a/tests/server/agent_execution/test_simple_request_context_builder.py b/tests/server/agent_execution/test_simple_request_context_builder.py index 5e1b8fd8..c1cbcf05 100644 --- a/tests/server/agent_execution/test_simple_request_context_builder.py +++ b/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()