fix: handle unavailable agents and retain fake manager state

Signed-off-by: Anas Khan <83116240+anxkhn@users.noreply.github.com>

Author: anxkhn

demo/ui/pages/agent_list.py

@@ -1,3 +1,4 @@
+# ruff: noqa
 import asyncio
 
 import mesop as me
@@ -20,13 +21,9 @@ def agent_list_page(app_state: AppState) -> None:
             with header('Remote Agents', 'smart_toy'):
                 pass
             agents = asyncio.run(ListRemoteAgents())
-            agents_list(agents)
+            agents_list(agents or [])
             with dialog(state.agent_dialog_open):
-                with me.box(
-                    style=me.Style(
-                        display='flex', flex_direction='column', gap=12
-                    )
-                ):
+                with me.box(style=me.Style(display='flex', flex_direction='column', gap=12)):
                     me.input(
                         label='Agent Address',
                         on_blur=set_agent_address,
@@ -43,9 +40,7 @@ def agent_list_page(app_state: AppState) -> None:
                     if state.agent_description:
                         me.text(f'Agent Description: {state.agent_description}')
                     if state.agent_framework_type:
-                        me.text(
-                            f'Agent Framework Type: {state.agent_framework_type}'
-                        )
+                        me.text(f'Agent Framework Type: {state.agent_framework_type}')
                     if state.input_modes:
                         me.text(f'Input Modes: {input_modes_string}')
                     if state.output_modes:
@@ -54,9 +49,7 @@ def agent_list_page(app_state: AppState) -> None:
                         me.text(f'Extensions: {extensions_string}')
 
                     if state.agent_name:
-                        me.text(
-                            f'Streaming Supported: {state.stream_supported}'
-                        )
+                        me.text(f'Streaming Supported: {state.stream_supported}')
                         me.text(
                             f'Push Notifications Supported: {state.push_notifications_supported}'
                         )
@@ -81,20 +74,14 @@ async def load_agent_info(e: me.ClickEvent) -> None:
         state.agent_name = agent_card_response.name
         state.agent_description = agent_card_response.description
         state.agent_framework_type = (
-            agent_card_response.provider.organization
-            if agent_card_response.provider
-            else ''
+            agent_card_response.provider.organization if agent_card_response.provider else ''
         )
         state.input_modes = agent_card_response.default_input_modes
         state.output_modes = agent_card_response.default_output_modes
         if agent_card_response.capabilities.extensions:
-            state.extensions = [
-                ext.uri for ext in agent_card_response.capabilities.extensions
-            ]
+            state.extensions = [ext.uri for ext in agent_card_response.capabilities.extensions]
         state.stream_supported = agent_card_response.capabilities.streaming
-        state.push_notifications_supported = (
-            agent_card_response.capabilities.push_notifications
-        )
+        state.push_notifications_supported = agent_card_response.capabilities.push_notifications
     except Exception as e:
         print(e)
         state.agent_name = None

demo/ui/service/server/in_memory_manager.py

@@ -1,3 +1,4 @@
+# ruff: noqa
 import asyncio
 import datetime
 import uuid
@@ -147,14 +148,10 @@ def add_event(self, event: Event):
 
     def next_message(self) -> Message:
         message = _message_queue[self._next_message_idx]
-        self._next_message_idx = (self._next_message_idx + 1) % len(
-            _message_queue
-        )
+        self._next_message_idx = (self._next_message_idx + 1) % len(_message_queue)
         return message
 
-    def get_conversation(
-        self, conversation_id: str | None
-    ) -> Conversation | None:
+    def get_conversation(self, conversation_id: str | None) -> Conversation | None:
         if not conversation_id:
             return None
         return next(
@@ -170,9 +167,7 @@ def get_pending_messages(self) -> list[tuple[str, str]]:
         for message_id in self._pending_message_ids:
             if message_id in self._task_map:
                 task_id = self._task_map[message_id]
-                task = next(
-                    filter(lambda x: x.id == task_id, self._tasks), None
-                )
+                task = next(filter(lambda x: x.id == task_id, self._tasks), None)
                 if not task:
                     rval.append((message_id, ''))
                 elif task.history and task.history[-1].parts:
@@ -183,15 +178,14 @@ def get_pending_messages(self) -> list[tuple[str, str]]:
                         rval.append(
                             (
                                 message_id,
-                                part.root.text
-                                if part.root.kind == 'text'
-                                else 'Working...',
+                                part.root.text if part.root.kind == 'text' else 'Working...',
                             )
                         )
+                else:
+                    rval.append((message_id, ''))
             else:
                 rval.append((message_id, ''))
-            return rval
-        return [(x, '') for x in self._pending_message_ids]
+        return rval
 
     def register_agent(self, url):
         agent_data = get_agent_card(url)
@@ -213,7 +207,7 @@ def tasks(self) -> list[Task]:
 
     @property
     def events(self) -> list[Event]:
-        return []
+        return self._events
 
 
 _contextId = str(uuid.uuid4())

demo/ui/state/host_agent_service.py

@@ -1,3 +1,4 @@
+# ruff: noqa
 import json
 import os
 import sys
@@ -75,9 +76,10 @@ async def ListRemoteAgents():
     client = ConversationClient(server_url)
     try:
         response = await client.list_agents(ListAgentRequest())
-        return response.result
+        return response.result if response.result else []
     except Exception as e:
         print('Failed to read agents', e)
+    return []
 
 
 async def AddRemoteAgent(path: str):
@@ -102,9 +104,10 @@ async def GetProcessingMessages():
     client = ConversationClient(server_url)
     try:
         response = await client.get_pending_messages(PendingMessageRequest())
-        return dict(response.result)
+        return dict(response.result) if response.result else {}
     except Exception as e:
         print('Error getting pending messages', e)
+    return {}
 
 
 def GetMessageAliases():
@@ -115,18 +118,16 @@ async def GetTasks():
     client = ConversationClient(server_url)
     try:
         response = await client.list_tasks(ListTaskRequest())
-        return response.result
+        return response.result if response.result else []
     except Exception as e:
         print('Failed to list tasks ', e)
-        return []
+    return []
 
 
 async def ListMessages(conversation_id: str) -> list[Message]:
     client = ConversationClient(server_url)
     try:
-        response = await client.list_messages(
-            ListMessageRequest(params=conversation_id)
-        )
+        response = await client.list_messages(ListMessageRequest(params=conversation_id))
         return response.result if response.result else []
     except Exception as e:
         print('Failed to list messages ', e)
@@ -147,9 +148,7 @@ async def UpdateAppState(state: AppState, conversation_id: str):
         if not conversations:
             state.conversations = []
         else:
-            state.conversations = [
-                convert_conversation_to_state(x) for x in conversations
-            ]
+            state.conversations = [convert_conversation_to_state(x) for x in conversations]
 
         state.task_list = []
         for task in await GetTasks():
@@ -176,9 +175,7 @@ async def UpdateApiKey(api_key: str):
 
         # Call the update API endpoint
         async with httpx.AsyncClient() as client:
-            response = await client.post(
-                f'{server_url}/api_key/update', json={'api_key': api_key}
-            )
+            response = await client.post(f'{server_url}/api_key/update', json={'api_key': api_key})
             response.raise_for_status()
         return True
     except Exception as e:
@@ -212,11 +209,7 @@ def convert_conversation_to_state(
 
 def convert_task_to_state(task: Task) -> StateTask:
     # Get the first message as the description
-    output = (
-        [extract_content(a.parts) for a in task.artifacts]
-        if task.artifacts
-        else []
-    )
+    output = [extract_content(a.parts) for a in task.artifacts] if task.artifacts else []
     if not task.history:
         return StateTask(
             task_id=task.id,