Merge branch 'main' into test-card-resolver

Author: lkawka

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/auth/test_user.py

@@ -1,9 +1,19 @@
 import unittest
 
-from a2a.auth.user import UnauthenticatedUser
+from inspect import isabstract
+
+from a2a.auth.user import UnauthenticatedUser, User
+
+
+class TestUser(unittest.TestCase):
+    def test_is_abstract(self):
+        self.assertTrue(isabstract(User))
 
 
 class TestUnauthenticatedUser(unittest.TestCase):
+    def test_is_user_subclass(self):
+        self.assertTrue(issubclass(UnauthenticatedUser, User))
+
     def test_is_authenticated_returns_false(self):
         user = UnauthenticatedUser()
         self.assertFalse(user.is_authenticated)

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()