Skip to content

Commit a7ed7ef

Browse files
fix: Add Input Validation for Task Context IDs in new_task Function (#340)
# Description This update addresses an issue where the `new_task` function in `src/a2a/utils/task.py` did not validate the format of the `context_id` when it was provided. Invalid or malformed `context_id` values, such as empty strings or non-UUID strings, were not being properly handled, which could lead to downstream errors. The changes in this PR introduce a validation check to ensure that any provided `context_id` is a valid UUID. If the `context_id` is invalid, a `ValueError` is raised. If no `context_id` is provided, a new UUID is generated as before. --------- Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
1 parent 4d5b92c commit a7ed7ef

File tree

2 files changed

+32
-4
lines changed

2 files changed

+32
-4
lines changed

src/a2a/utils/task.py

Lines changed: 14 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,7 @@ def new_task(request: Message) -> Task:
1818
1919
Raises:
2020
TypeError: If the message role is None.
21-
ValueError: If the message parts are empty or if any part has empty content.
21+
ValueError: If the message parts are empty, if any part has empty content, or if the provided context_id is invalid.
2222
"""
2323
if not request.role:
2424
raise TypeError('Message role cannot be None')
@@ -28,12 +28,22 @@ def new_task(request: Message) -> Task:
2828
if isinstance(part.root, TextPart) and not part.root.text:
2929
raise ValueError('TextPart content cannot be empty')
3030

31+
context_id_str = request.context_id
32+
if context_id_str is not None:
33+
try:
34+
uuid.UUID(context_id_str)
35+
context_id = context_id_str
36+
except (ValueError, AttributeError, TypeError) as e:
37+
raise ValueError(
38+
f"Invalid context_id: '{context_id_str}' is not a valid UUID."
39+
) from e
40+
else:
41+
context_id = str(uuid.uuid4())
42+
3143
return Task(
3244
status=TaskStatus(state=TaskState.submitted),
3345
id=(request.task_id if request.task_id else str(uuid.uuid4())),
34-
context_id=(
35-
request.context_id if request.context_id else str(uuid.uuid4())
36-
),
46+
context_id=context_id,
3747
history=[request],
3848
)
3949

tests/utils/test_task.py

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -188,6 +188,24 @@ def test_completed_task_invalid_artifact_type(self):
188188
history=[],
189189
)
190190

191+
def test_new_task_with_invalid_context_id(self):
192+
"""Test that new_task raises a ValueError for various invalid context_id formats."""
193+
invalid_ids = ['not-a-uuid', '']
194+
for invalid_id in invalid_ids:
195+
with self.subTest(invalid_id=invalid_id):
196+
with pytest.raises(
197+
ValueError,
198+
match=f"Invalid context_id: '{invalid_id}' is not a valid UUID.",
199+
):
200+
new_task(
201+
Message(
202+
role=Role.user,
203+
parts=[Part(root=TextPart(text='test message'))],
204+
message_id=str(uuid.uuid4()),
205+
context_id=invalid_id,
206+
)
207+
)
208+
191209

192210
if __name__ == '__main__':
193211
unittest.main()

0 commit comments

Comments
 (0)