From a7800a03cbdb9f6249a8c9277ca6620a9826d6b0 Mon Sep 17 00:00:00 2001 From: Pankaj Kumar Bind Date: Thu, 24 Jul 2025 00:46:37 +0530 Subject: [PATCH 1/6] feat: Add Input Validation for Task Context IDs in new_task Function --- src/a2a/utils/task.py | 21 ++++++++++++++++---- tests/utils/test_task.py | 42 ++++++++++++++++++++++++++++++++++++++++ 2 files changed, 59 insertions(+), 4 deletions(-) diff --git a/src/a2a/utils/task.py b/src/a2a/utils/task.py index b57da9b8..7bc23ee7 100644 --- a/src/a2a/utils/task.py +++ b/src/a2a/utils/task.py @@ -18,7 +18,7 @@ def new_task(request: Message) -> Task: Raises: TypeError: If the message role is None. - ValueError: If the message parts are empty or if any part has empty content. + ValueError: If the message parts are empty, if any part has empty content, or if the provided context_id is invalid. """ if not request.role: raise TypeError('Message role cannot be None') @@ -28,12 +28,25 @@ def new_task(request: Message) -> Task: if isinstance(part.root, TextPart) and not part.root.text: raise ValueError('TextPart content cannot be empty') + context_id_str = request.context_id + if context_id_str is not None: + try: + # Validate that the provided context_id is a valid UUID + uuid.UUID(context_id_str) + context_id = context_id_str + except (ValueError, AttributeError, TypeError): + # Catch a variety of potential issues with the UUID validation + raise ValueError( + f"Invalid context_id: '{context_id_str}' is not a valid UUID." + ) + else: + # Generate a new UUID if no context_id is provided + context_id = str(uuid.uuid4()) + return Task( status=TaskStatus(state=TaskState.submitted), id=(request.task_id if request.task_id else str(uuid.uuid4())), - context_id=( - request.context_id if request.context_id else str(uuid.uuid4()) - ), + context_id=context_id, history=[request], ) diff --git a/tests/utils/test_task.py b/tests/utils/test_task.py index cb3dc386..8386d836 100644 --- a/tests/utils/test_task.py +++ b/tests/utils/test_task.py @@ -188,6 +188,48 @@ def test_completed_task_invalid_artifact_type(self): history=[], ) + def test_new_task_with_invalid_context_id(self): + """Test that new_task raises a ValueError with an invalid context_id.""" + with pytest.raises( + ValueError, + match="Invalid context_id: 'not-a-uuid' is not a valid UUID.", + ): + new_task( + Message( + role=Role.user, + parts=[Part(root=TextPart(text='test message'))], + message_id=str(uuid.uuid4()), + context_id='not-a-uuid', + ) + ) + + def test_new_task_with_empty_string_context_id(self): + """Test that new_task raises a ValueError with an empty string context_id.""" + with pytest.raises( + ValueError, match="Invalid context_id: '' is not a valid UUID." + ): + new_task( + Message( + role=Role.user, + parts=[Part(root=TextPart(text='test message'))], + message_id=str(uuid.uuid4()), + context_id='', + ) + ) + + def test_new_task_with_valid_context_id(self): + """Test that new_task accepts a valid context_id.""" + valid_uuid = '123e4567-e89b-12d3-a456-426614174000' + task = new_task( + Message( + role=Role.user, + parts=[Part(root=TextPart(text='test message'))], + message_id=str(uuid.uuid4()), + context_id=valid_uuid, + ) + ) + self.assertEqual(task.context_id, valid_uuid) + if __name__ == '__main__': unittest.main() From e1f2bdd5337586f89150486271bb7104b7aeb905 Mon Sep 17 00:00:00 2001 From: Pankaj Kumar Bind Date: Thu, 24 Jul 2025 00:49:33 +0530 Subject: [PATCH 2/6] fix test_task.py --- tests/utils/test_task.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/tests/utils/test_task.py b/tests/utils/test_task.py index 8386d836..7c207bbe 100644 --- a/tests/utils/test_task.py +++ b/tests/utils/test_task.py @@ -189,7 +189,6 @@ def test_completed_task_invalid_artifact_type(self): ) def test_new_task_with_invalid_context_id(self): - """Test that new_task raises a ValueError with an invalid context_id.""" with pytest.raises( ValueError, match="Invalid context_id: 'not-a-uuid' is not a valid UUID.", @@ -204,7 +203,6 @@ def test_new_task_with_invalid_context_id(self): ) def test_new_task_with_empty_string_context_id(self): - """Test that new_task raises a ValueError with an empty string context_id.""" with pytest.raises( ValueError, match="Invalid context_id: '' is not a valid UUID." ): @@ -218,7 +216,6 @@ def test_new_task_with_empty_string_context_id(self): ) def test_new_task_with_valid_context_id(self): - """Test that new_task accepts a valid context_id.""" valid_uuid = '123e4567-e89b-12d3-a456-426614174000' task = new_task( Message( From 00c3312014d12355bd5a6d69fa9c8d915ae04bae Mon Sep 17 00:00:00 2001 From: Pankaj Kumar Bind <73558583+pankaj-bind@users.noreply.github.com> Date: Thu, 24 Jul 2025 00:55:45 +0530 Subject: [PATCH 3/6] Update src/a2a/utils/task.py Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --- src/a2a/utils/task.py | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/src/a2a/utils/task.py b/src/a2a/utils/task.py index 7bc23ee7..cee28b85 100644 --- a/src/a2a/utils/task.py +++ b/src/a2a/utils/task.py @@ -29,18 +29,18 @@ def new_task(request: Message) -> Task: raise ValueError('TextPart content cannot be empty') context_id_str = request.context_id - if context_id_str is not None: + context_id = request.context_id + if context_id is not None: try: - # Validate that the provided context_id is a valid UUID - uuid.UUID(context_id_str) - context_id = context_id_str - except (ValueError, AttributeError, TypeError): - # Catch a variety of potential issues with the UUID validation + # Validate that the provided context_id is a valid UUID. + uuid.UUID(context_id) + except ValueError: + # Re-raise as ValueError with a more specific message. raise ValueError( - f"Invalid context_id: '{context_id_str}' is not a valid UUID." - ) + f"Invalid context_id: '{context_id}' is not a valid UUID." + ) from None else: - # Generate a new UUID if no context_id is provided + # Generate a new UUID if no context_id is provided. context_id = str(uuid.uuid4()) return Task( From eae5ac72474f5b3d5b9d391af0c2b960cfa3d718 Mon Sep 17 00:00:00 2001 From: Pankaj Kumar Bind Date: Thu, 24 Jul 2025 01:10:53 +0530 Subject: [PATCH 4/6] fix linting --- src/a2a/utils/task.py | 15 ++++++--------- 1 file changed, 6 insertions(+), 9 deletions(-) diff --git a/src/a2a/utils/task.py b/src/a2a/utils/task.py index cee28b85..60272367 100644 --- a/src/a2a/utils/task.py +++ b/src/a2a/utils/task.py @@ -29,18 +29,15 @@ def new_task(request: Message) -> Task: raise ValueError('TextPart content cannot be empty') context_id_str = request.context_id - context_id = request.context_id - if context_id is not None: + if context_id_str is not None: try: - # Validate that the provided context_id is a valid UUID. - uuid.UUID(context_id) - except ValueError: - # Re-raise as ValueError with a more specific message. + uuid.UUID(context_id_str) + context_id = context_id_str + except (ValueError, AttributeError, TypeError) as e: raise ValueError( - f"Invalid context_id: '{context_id}' is not a valid UUID." - ) from None + f"Invalid context_id: '{context_id_str}' is not a valid UUID." + ) from e else: - # Generate a new UUID if no context_id is provided. context_id = str(uuid.uuid4()) return Task( From f4d63a31c963f650caa050709319a3e28b6e30c2 Mon Sep 17 00:00:00 2001 From: Pankaj Kumar Bind Date: Thu, 24 Jul 2025 01:23:49 +0530 Subject: [PATCH 5/6] merge tests --- tests/utils/test_task.py | 53 ++++++++++++---------------------------- 1 file changed, 16 insertions(+), 37 deletions(-) diff --git a/tests/utils/test_task.py b/tests/utils/test_task.py index 7c207bbe..5cc35229 100644 --- a/tests/utils/test_task.py +++ b/tests/utils/test_task.py @@ -189,43 +189,22 @@ def test_completed_task_invalid_artifact_type(self): ) def test_new_task_with_invalid_context_id(self): - with pytest.raises( - ValueError, - match="Invalid context_id: 'not-a-uuid' is not a valid UUID.", - ): - new_task( - Message( - role=Role.user, - parts=[Part(root=TextPart(text='test message'))], - message_id=str(uuid.uuid4()), - context_id='not-a-uuid', - ) - ) - - def test_new_task_with_empty_string_context_id(self): - with pytest.raises( - ValueError, match="Invalid context_id: '' is not a valid UUID." - ): - new_task( - Message( - role=Role.user, - parts=[Part(root=TextPart(text='test message'))], - message_id=str(uuid.uuid4()), - context_id='', - ) - ) - - def test_new_task_with_valid_context_id(self): - valid_uuid = '123e4567-e89b-12d3-a456-426614174000' - task = new_task( - Message( - role=Role.user, - parts=[Part(root=TextPart(text='test message'))], - message_id=str(uuid.uuid4()), - context_id=valid_uuid, - ) - ) - self.assertEqual(task.context_id, valid_uuid) + """Test that new_task raises a ValueError for various invalid context_id formats.""" + invalid_ids = ["not-a-uuid", ""] + for invalid_id in invalid_ids: + with self.subTest(invalid_id=invalid_id): + with pytest.raises( + ValueError, + match=f"Invalid context_id: '{invalid_id}' is not a valid UUID.", + ): + new_task( + Message( + role=Role.user, + parts=[Part(root=TextPart(text='test message'))], + message_id=str(uuid.uuid4()), + context_id=invalid_id, + ) + ) if __name__ == '__main__': From d07a7667d0a4a39c347d2c4befc051ba6167a131 Mon Sep 17 00:00:00 2001 From: Pankaj Kumar Bind Date: Thu, 24 Jul 2025 01:26:02 +0530 Subject: [PATCH 6/6] fix linting --- tests/utils/test_task.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/utils/test_task.py b/tests/utils/test_task.py index 5cc35229..77441316 100644 --- a/tests/utils/test_task.py +++ b/tests/utils/test_task.py @@ -190,7 +190,7 @@ def test_completed_task_invalid_artifact_type(self): def test_new_task_with_invalid_context_id(self): """Test that new_task raises a ValueError for various invalid context_id formats.""" - invalid_ids = ["not-a-uuid", ""] + invalid_ids = ['not-a-uuid', ''] for invalid_id in invalid_ids: with self.subTest(invalid_id=invalid_id): with pytest.raises(