Fix sut_agent.py

Author: ishymko

tck/sut_agent.py

@@ -18,12 +18,14 @@
 from a2a.types import (
     AgentCapabilities,
     AgentCard,
+    AgentInterface,
     AgentProvider,
+    AgentSkill,
     Message,
+    Part,
     TaskState,
     TaskStatus,
     TaskStatusUpdateEvent,
-    TextPart,
 )
 
 
@@ -38,24 +40,25 @@ class SUTAgentExecutor(AgentExecutor):
 
     def __init__(self) -> None:
         """Initializes the SUT agent executor."""
-        self.running_tasks = set()
+        self.running_tasks: set[str] = set()
 
     async def cancel(
         self, context: RequestContext, event_queue: EventQueue
     ) -> None:
         """Cancels a task."""
         api_task_id = context.task_id
+        if api_task_id is None:
+            return
         if api_task_id in self.running_tasks:
             self.running_tasks.remove(api_task_id)
 
         status_update = TaskStatusUpdateEvent(
             task_id=api_task_id,
             context_id=context.context_id or str(uuid.uuid4()),
             status=TaskStatus(
-                state=TaskState.canceled,
-                timestamp=datetime.now(timezone.utc).isoformat(),
+                state=TaskState.TASK_STATE_CANCELED,
+                timestamp=datetime.now(timezone.utc),
             ),
-            final=True,
         )
         await event_queue.enqueue_event(status_update)
 
@@ -65,6 +68,8 @@ async def execute(
         """Executes a task."""
         user_message = context.message
         task_id = context.task_id
+        if user_message is None or task_id is None:
+            return
         context_id = context.context_id
 
         self.running_tasks.add(task_id)
@@ -80,17 +85,16 @@ async def execute(
             task_id=task_id,
             context_id=context_id,
             status=TaskStatus(
-                state=TaskState.working,
+                state=TaskState.TASK_STATE_WORKING,
                 message=Message(
                     role='agent',
                     message_id=str(uuid.uuid4()),
-                    parts=[TextPart(text='Processing your question')],
+                    parts=[Part(text='Processing your question')],
                     task_id=task_id,
                     context_id=context_id,
                 ),
-                timestamp=datetime.now(timezone.utc).isoformat(),
+                timestamp=datetime.now(timezone.utc),
             ),
-            final=False,
         )
         await event_queue.enqueue_event(working_status)
 
@@ -106,7 +110,7 @@ async def execute(
         agent_message = Message(
             role='agent',
             message_id=str(uuid.uuid4()),
-            parts=[TextPart(text=agent_reply_text)],
+            parts=[Part(text=agent_reply_text)],
             task_id=task_id,
             context_id=context_id,
         )
@@ -115,11 +119,10 @@ async def execute(
             task_id=task_id,
             context_id=context_id,
             status=TaskStatus(
-                state=TaskState.input_required,
+                state=TaskState.TASK_STATE_INPUT_REQUIRED,
                 message=agent_message,
-                timestamp=datetime.now(timezone.utc).isoformat(),
+                timestamp=datetime.now(timezone.utc),
             ),
-            final=True,
         )
         await event_queue.enqueue_event(final_update)
 
@@ -131,38 +134,34 @@ def main() -> None:
     agent_card = AgentCard(
         name='SUT Agent',
         description='An agent to be used as SUT against TCK tests.',
-        url=f'http://localhost:{http_port}{JSONRPC_URL}',
+        supported_interfaces=[
+            AgentInterface(
+                url=f'http://localhost:{http_port}{JSONRPC_URL}',
+                protocol_binding='JSONRPC',
+                protocol_version='0.3.0',
+            ),
+        ],
         provider=AgentProvider(
             organization='A2A Samples',
             url='https://example.com/a2a-samples',
         ),
         version='1.0.0',
-        protocol_version='0.3.0',
         capabilities=AgentCapabilities(
             streaming=True,
             push_notifications=False,
-            state_transition_history=True,
         ),
         default_input_modes=['text'],
         default_output_modes=['text', 'task-status'],
         skills=[
-            {
-                'id': 'sut_agent',
-                'name': 'SUT Agent',
-                'description': 'Simulate the general flow of a streaming agent.',
-                'tags': ['sut'],
-                'examples': ['hi', 'hello world', 'how are you', 'goodbye'],
-                'input_modes': ['text'],
-                'output_modes': ['text', 'task-status'],
-            }
-        ],
-        supports_authenticated_extended_card=False,
-        preferred_transport='JSONRPC',
-        additional_interfaces=[
-            {
-                'url': f'http://localhost:{http_port}{JSONRPC_URL}',
-                'transport': 'JSONRPC',
-            },
+            AgentSkill(
+                id='sut_agent',
+                name='SUT Agent',
+                description='Simulate the general flow of a streaming agent.',
+                tags=['sut'],
+                examples=['hi', 'hello world', 'how are you', 'goodbye'],
+                input_modes=['text'],
+                output_modes=['text', 'task-status'],
+            )
         ],
     )