diff --git a/src/a2a/server/agent_execution/active_task_registry.py b/src/a2a/server/agent_execution/active_task_registry.py index 9c1299ab3..4f5822f68 100644 --- a/src/a2a/server/agent_execution/active_task_registry.py +++ b/src/a2a/server/agent_execution/active_task_registry.py @@ -11,6 +11,7 @@ from a2a.server.context import ServerCallContext from a2a.server.tasks.push_notification_sender import PushNotificationSender from a2a.server.tasks.task_store import TaskStore + from a2a.types.a2a_pb2 import Message from a2a.server.agent_execution.active_task import ActiveTask from a2a.server.tasks.task_manager import TaskManager @@ -41,6 +42,7 @@ async def get_or_create( call_context: ServerCallContext, context_id: str | None = None, create_task_if_missing: bool = False, + initial_message: Message | None = None, ) -> ActiveTask: """Retrieves an existing ActiveTask or creates a new one.""" async with self._lock: @@ -51,7 +53,7 @@ async def get_or_create( task_id=task_id, context_id=context_id, task_store=self._task_store, - initial_message=None, + initial_message=initial_message, context=call_context, ) diff --git a/src/a2a/server/request_handlers/default_request_handler_v2.py b/src/a2a/server/request_handlers/default_request_handler_v2.py index 30304609a..28b60a20c 100644 --- a/src/a2a/server/request_handlers/default_request_handler_v2.py +++ b/src/a2a/server/request_handlers/default_request_handler_v2.py @@ -222,6 +222,7 @@ async def _setup_active_task( context_id=context_id, call_context=call_context, create_task_if_missing=True, + initial_message=params.message, ) return active_task, request_context