From f666bec5417860745874025811fb130a2fda158d Mon Sep 17 00:00:00 2001 From: Will Chen Date: Mon, 15 Dec 2025 12:24:05 -0800 Subject: [PATCH 1/6] feat: Add custom id generators option to simple_request_context_builder --- .../agent_execution/simple_request_context_builder.py | 9 +++++++++ 1 file changed, 9 insertions(+) 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..c2f68327 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, ) From 98d4414b3ab8ca70089a060cf1773c3ccbc94b29 Mon Sep 17 00:00:00 2001 From: Will Chen Date: Mon, 15 Dec 2025 12:53:21 -0800 Subject: [PATCH 2/6] feat: Add custom id generators option to simple_request_context_builder --- .../agent_execution/simple_request_context_builder.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) 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 c2f68327..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,7 +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.id_generator import IDGenerator from a2a.server.tasks import TaskStore from a2a.types import MessageSendParams, Task @@ -14,8 +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, + task_id_generator: IDGenerator | None = None, + context_id_generator: IDGenerator | None = None, ) -> None: """Initializes the SimpleRequestContextBuilder. From 3a45f32598c0aa0498c8bd88159363e405a1f52f Mon Sep 17 00:00:00 2001 From: Will Chen Date: Mon, 15 Dec 2025 16:03:19 -0800 Subject: [PATCH 3/6] Add tests --- .../test_simple_request_context_builder.py | 54 +++++++++++++++++++ 1 file changed, 54 insertions(+) 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..d6e1b729 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, @@ -276,5 +277,58 @@ async def test_build_populate_false_with_reference_task_ids(self) -> None: 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) + + 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() + + 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() From e4ee5f26ae9ba6d52f7a05a360594083c2fcd6f4 Mon Sep 17 00:00:00 2001 From: Will Chen Date: Mon, 15 Dec 2025 16:13:04 -0800 Subject: [PATCH 4/6] test: Add tests for custom id generators --- .../agent_execution/test_simple_request_context_builder.py | 4 ++++ 1 file changed, 4 insertions(+) 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 d6e1b729..587c48ae 100644 --- a/tests/server/agent_execution/test_simple_request_context_builder.py +++ b/tests/server/agent_execution/test_simple_request_context_builder.py @@ -280,6 +280,8 @@ async def test_build_populate_false_with_reference_task_ids(self) -> None: 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, @@ -300,6 +302,8 @@ async def test_build_with_custom_id_generators(self) -> None: 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) From fffc9ead1436e251d0d11ddd2a44db3861f5c6b5 Mon Sep 17 00:00:00 2001 From: Will Chen Date: Mon, 15 Dec 2025 16:21:01 -0800 Subject: [PATCH 5/6] Fix lint --- .../agent_execution/test_simple_request_context_builder.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) 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 587c48ae..63e2e842 100644 --- a/tests/server/agent_execution/test_simple_request_context_builder.py +++ b/tests/server/agent_execution/test_simple_request_context_builder.py @@ -276,7 +276,6 @@ 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) @@ -299,7 +298,7 @@ async def test_build_with_custom_id_generators(self) -> 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') @@ -328,7 +327,7 @@ async def test_build_with_provided_ids_and_custom_id_generators(self) -> None: 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) From 774fed016fd84baa30e0894d72a841acee2f0430 Mon Sep 17 00:00:00 2001 From: Will Chen Date: Mon, 15 Dec 2025 16:52:05 -0800 Subject: [PATCH 6/6] fix lint --- .../agent_execution/test_simple_request_context_builder.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) 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 63e2e842..c1cbcf05 100644 --- a/tests/server/agent_execution/test_simple_request_context_builder.py +++ b/tests/server/agent_execution/test_simple_request_context_builder.py @@ -304,7 +304,9 @@ async def test_build_with_custom_id_generators(self) -> None: 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: + 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) @@ -333,5 +335,6 @@ async def test_build_with_provided_ids_and_custom_id_generators(self) -> None: self.assertEqual(request_context.task_id, provided_task_id) self.assertEqual(request_context.context_id, provided_context_id) + if __name__ == '__main__': unittest.main()