Add input required test

Author: ishymko

tests/integration/test_end_to_end.py

@@ -39,16 +39,45 @@ async def execute(self, context: RequestContext, event_queue: EventQueue):
             context.task_id,
             context.context_id,
         )
-        await task_updater.update_status(TaskState.TASK_STATE_SUBMITTED)
-        await task_updater.update_status(TaskState.TASK_STATE_WORKING)
-        await task_updater.add_artifact(
-            parts=[Part(text='artifact content')], name='test-artifact'
+        user_input = context.get_user_input()
+
+        is_input_required_resumption = (
+            context.current_task is not None
+            and context.current_task.status.state
+            == TaskState.TASK_STATE_INPUT_REQUIRED
         )
+
+        if not is_input_required_resumption:
+            await task_updater.update_status(
+                TaskState.TASK_STATE_SUBMITTED,
+                message=task_updater.new_agent_message(
+                    [Part(text='task submitted')]
+                ),
+            )
+
         await task_updater.update_status(
-            TaskState.TASK_STATE_COMPLETED,
-            message=task_updater.new_agent_message([Part(text='done')]),
+            TaskState.TASK_STATE_WORKING,
+            message=task_updater.new_agent_message(
+                [Part(text='task working')]
+            ),
         )
 
+        if user_input == 'Need input':
+            await task_updater.update_status(
+                TaskState.TASK_STATE_INPUT_REQUIRED,
+                message=task_updater.new_agent_message(
+                    [Part(text='Please provide input')]
+                ),
+            )
+        else:
+            await task_updater.add_artifact(
+                parts=[Part(text='artifact content')], name='test-artifact'
+            )
+            await task_updater.update_status(
+                TaskState.TASK_STATE_COMPLETED,
+                message=task_updater.new_agent_message([Part(text='done')]),
+            )
+
     async def cancel(self, context: RequestContext, event_queue: EventQueue):
         raise NotImplementedError('Cancellation is not supported')
 
@@ -221,9 +250,13 @@ async def test_end_to_end_send_message_blocking(transport_setups):
     assert len(response.task.artifacts) == 1
     assert response.task.artifacts[0].name == 'test-artifact'
     assert response.task.artifacts[0].parts[0].text == 'artifact content'
-    assert len(response.task.history) == 1
+    assert len(response.task.history) == 3
     assert response.task.history[0].role == Role.ROLE_USER
     assert response.task.history[0].parts[0].text == 'Run dummy agent!'
+    assert response.task.history[1].role == Role.ROLE_AGENT
+    assert response.task.history[1].parts[0].text == 'task submitted'
+    assert response.task.history[2].role == Role.ROLE_AGENT
+    assert response.task.history[2].parts[0].text == 'task working'
 
 
 @pytest.mark.asyncio
@@ -275,7 +308,7 @@ async def test_end_to_end_send_message_streaming(transport_setups):
     ]
 
     assert len(events) == len(expected_events)
-    for (event, task), (expected_type, expected_state) in zip(
+    for (event, _), (expected_type, expected_state) in zip(
         events, expected_events, strict=True
     ):
         assert event.HasField(expected_type)
@@ -289,9 +322,13 @@ async def test_end_to_end_send_message_streaming(transport_setups):
             )
 
     last_task = events[-1][1]
-    assert len(last_task.history) == 1
+    assert len(last_task.history) == 3
     assert last_task.history[0].role == Role.ROLE_AGENT
-    assert last_task.history[0].parts[0].text == 'done'
+    assert last_task.history[0].parts[0].text == 'task submitted'
+    assert last_task.history[1].role == Role.ROLE_AGENT
+    assert last_task.history[1].parts[0].text == 'task working'
+    assert last_task.history[2].role == Role.ROLE_AGENT
+    assert last_task.history[2].parts[0].text == 'done'
 
 
 @pytest.mark.asyncio
@@ -318,9 +355,13 @@ async def test_end_to_end_get_task(transport_setups):
         TaskState.TASK_STATE_WORKING,
         TaskState.TASK_STATE_COMPLETED,
     }
-    assert len(retrieved_task.history) == 1
+    assert len(retrieved_task.history) == 3
     assert retrieved_task.history[0].role == Role.ROLE_USER
     assert retrieved_task.history[0].parts[0].text == 'Test Get Task'
+    assert retrieved_task.history[1].role == Role.ROLE_AGENT
+    assert retrieved_task.history[1].parts[0].text == 'task submitted'
+    assert retrieved_task.history[2].role == Role.ROLE_AGENT
+    assert retrieved_task.history[2].parts[0].text == 'task working'
 
 
 @pytest.mark.asyncio
@@ -361,11 +402,104 @@ async def test_end_to_end_list_tasks(transport_setups):
         actual_task_ids.extend([task.id for task in list_response.tasks])
 
         for task in list_response.tasks:
-            assert len(task.history) == 1
+            assert len(task.history) >= 1
             assert task.history[0].role == Role.ROLE_USER
             assert task.history[0].parts[0].text.startswith('Test List Tasks ')
 
         token = list_response.next_page_token
 
     assert len(actual_task_ids) == total_items
     assert sorted(actual_task_ids) == sorted(expected_task_ids)
+
+
+@pytest.mark.asyncio
+async def test_end_to_end_input_required(transport_setups):
+    client = transport_setups.client
+
+    message_to_send = Message(
+        role=Role.ROLE_USER,
+        message_id='msg-e2e-input-req-1',
+        parts=[Part(text='Need input')],
+    )
+
+    events = [
+        event async for event in client.send_message(request=message_to_send)
+    ]
+
+    expected_first_events = [
+        ('status_update', TaskState.TASK_STATE_SUBMITTED),
+        ('status_update', TaskState.TASK_STATE_WORKING),
+        ('status_update', TaskState.TASK_STATE_INPUT_REQUIRED),
+    ]
+
+    assert len(events) == len(expected_first_events)
+    for (event, _), (expected_type, expected_state) in zip(
+        events, expected_first_events, strict=True
+    ):
+        assert event.HasField(expected_type)
+        if expected_type == 'status_update':
+            assert event.status_update.status.state == expected_state
+
+    _, task = events[-1]
+
+    assert task.status.state == TaskState.TASK_STATE_INPUT_REQUIRED
+    assert len(task.history) == 3
+    assert task.history[0].role == Role.ROLE_AGENT
+    assert task.history[0].parts[0].text == 'task submitted'
+    assert task.history[1].role == Role.ROLE_AGENT
+    assert task.history[1].parts[0].text == 'task working'
+    assert task.history[2].role == Role.ROLE_AGENT
+    assert task.history[2].parts[0].text == 'Please provide input'
+    assert task.status.message.role == Role.ROLE_AGENT
+    assert task.status.message.parts[0].text == 'Please provide input'
+
+    task_id = task.id
+
+    # Follow-up message
+    follow_up_message = Message(
+        task_id=task_id,
+        role=Role.ROLE_USER,
+        message_id='msg-e2e-input-req-2',
+        parts=[Part(text='Here is the input')],
+    )
+
+    follow_up_events = [
+        event async for event in client.send_message(request=follow_up_message)
+    ]
+
+    expected_second_events = [
+        ('status_update', TaskState.TASK_STATE_WORKING),
+        ('artifact_update', None),
+        ('status_update', TaskState.TASK_STATE_COMPLETED),
+    ]
+
+    assert len(follow_up_events) == len(expected_second_events)
+    for (event, _), (expected_type, expected_state) in zip(
+        follow_up_events, expected_second_events, strict=True
+    ):
+        assert event.HasField(expected_type)
+        if expected_type == 'status_update':
+            assert event.status_update.status.state == expected_state
+        elif expected_type == 'artifact_update':
+            assert event.artifact_update.artifact.name == 'test-artifact'
+            assert (
+                event.artifact_update.artifact.parts[0].text
+                == 'artifact content'
+            )
+
+    _, task_end = follow_up_events[-1]
+    print(task_end)
+
+    assert task_end.status.state == TaskState.TASK_STATE_COMPLETED
+    assert task_end.id == task_id
+    assert len(task_end.artifacts) == 1
+    assert task_end.artifacts[0].name == 'test-artifact'
+    assert task_end.artifacts[0].parts[0].text == 'artifact content'
+
+    assert len(task_end.history) == 2
+    assert task_end.history[0].role == Role.ROLE_AGENT
+    assert task_end.history[0].parts[0].text == 'task working'
+    assert task_end.history[1].role == Role.ROLE_AGENT
+    assert task_end.history[1].parts[0].text == 'done'
+    assert task_end.status.message.role == Role.ROLE_AGENT
+    assert task_end.status.message.parts[0].text == 'done'