Merge branch 'jscpd-fixes' of https://github.com/google/a2a-python into jscpd-fixes

Author: holtskinner

examples/google_adk/calendar_agent/adk_agent_executor.py

@@ -9,6 +9,9 @@
 from google.adk import Runner
 from google.adk.auth import AuthConfig
 from google.adk.events import Event
+from google.adk.session_service import (
+    Session,
+)
 from google.genai import types
 
 from a2a.server.agent_execution import AgentExecutor
@@ -45,6 +48,7 @@ class ADKAgentExecutor(AgentExecutor):
     """An AgentExecutor that runs an ADK-based Agent."""
 
     _awaiting_auth: dict[str, asyncio.Future]
+    _running_sessions: dict[str, Session]
 
     def __init__(self, runner: Runner, card: AgentCard):
         self.runner = runner
@@ -98,12 +102,17 @@ async def _process_request(
                 # is received.
                 break
             if event.is_final_response():
-                parts = convert_genai_parts_to_a2a(event.content.parts)
-                logger.debug('Yielding final response: %s', parts)
-                task_updater.add_artifact(parts)
-                task_updater.complete()
+                if event.content and event.content.parts:
+                    parts = convert_genai_parts_to_a2a(event.content.parts)
+                    logger.debug('Yielding final response: %s', parts)
+                    task_updater.add_artifact(parts)
+                    task_updater.complete()
                 break
-            if not event.get_function_calls():
+            if (
+                not event.get_function_calls()
+                and event.content
+                and event.content.parts
+            ):
                 logger.debug('Yielding update response')
                 task_updater.update_status(
                     TaskState.working,
@@ -208,11 +217,29 @@ async def execute(
         event_queue: EventQueue,
     ):
         # Run the agent until either complete or the task is suspended.
+        assert context.task_id is not None, (
+            'Task ID must be present for execution'
+        )
+        assert context.context_id is not None, (
+            'Context ID must be present for execution'
+        )
+
         updater = TaskUpdater(event_queue, context.task_id, context.context_id)
         # Immediately notify that the task is submitted.
         if not context.current_task:
             updater.submit()
         updater.start_work()
+
+        if context.message is None:
+            logger.warning('Execute called with no message in context')
+            updater.update_status(
+                TaskState.failed,
+                message=new_agent_text_message(
+                    'No message provided to the agent.'
+                ),
+            )
+            return
+
         await self._process_request(
             types.UserContent(
                 parts=convert_a2a_parts_to_genai(context.message.parts),
@@ -229,12 +256,17 @@ async def cancel(self, context: RequestContext, event_queue: EventQueue):
     async def on_auth_callback(self, state: str, uri: str):
         self._awaiting_auth[state].set_result(uri)
 
-    def _upsert_session(self, session_id: str):
-        return self.runner.session_service.get_session(
-            app_name=self.runner.app_name, user_id='self', session_id=session_id
-        ) or self.runner.session_service.create_session(
+    def _upsert_session(self, session_id: str) -> Session:
+        session = self.runner.session_service.get_session(
             app_name=self.runner.app_name, user_id='self', session_id=session_id
         )
+        if session is None:
+            session = self.runner.session_service.create_session(
+                app_name=self.runner.app_name,
+                user_id='self',
+                session_id=session_id,
+            )
+        return session
 
 
 def convert_a2a_parts_to_genai(parts: list[Part]) -> list[types.Part]:
@@ -244,32 +276,43 @@ def convert_a2a_parts_to_genai(parts: list[Part]) -> list[types.Part]:
 
 def convert_a2a_part_to_genai(part: Part) -> types.Part:
     """Convert a single A2A Part type into a Google Gen AI Part type."""
-    part = part.root
-    if isinstance(part, TextPart):
-        return types.Part(text=part.text)
-    if isinstance(part, FilePart):
-        if isinstance(part.file, FileWithUri):
+    if isinstance(part.root, TextPart):
+        return types.Part(text=part.root.text)
+    if isinstance(part.root, FilePart):
+        file_data = part.root.file
+        if isinstance(file_data, FileWithUri):
             return types.Part(
                 file_data=types.FileData(
-                    file_uri=part.file.uri, mime_type=part.file.mime_type
+                    file_uri=file_data.uri,
+                    mime_type=file_data.mime_type,
                 )
             )
-        if isinstance(part.file, FileWithBytes):
+        if isinstance(file_data, FileWithBytes):
             return types.Part(
                 inline_data=types.Blob(
-                    data=part.file.bytes, mime_type=part.file.mime_type
+                    data=file_data.bytes,
+                    mime_type=file_data.mime_type,
                 )
             )
-        raise ValueError(f'Unsupported file type: {type(part.file)}')
-    raise ValueError(f'Unsupported part type: {type(part)}')
+        raise ValueError(f'Unsupported file type: {type(file_data)}')
+    raise ValueError(f'Unsupported part root type: {type(part.root)}')
 
 
 def convert_genai_parts_to_a2a(parts: list[types.Part]) -> list[Part]:
     """Convert a list of Google Gen AI Part types into a list of A2A Part types."""
+    if not parts:
+        return []
     return [
         convert_genai_part_to_a2a(part)
         for part in parts
-        if (part.text or part.file_data or part.inline_data)
+        if part
+        and (
+            part.text
+            or part.file_data
+            or part.inline_data
+            or part.function_call
+            or part.function_response
+        )
     ]
 
 
@@ -296,7 +339,9 @@ def convert_genai_part_to_a2a(part: types.Part) -> Part:
     raise ValueError(f'Unsupported part type: {part}')
 
 
-def get_auth_request_function_call(event: Event) -> types.FunctionCall:
+def get_auth_request_function_call(
+    event: Event,
+) -> types.FunctionCall | None:
     """Get the special auth request function call from the event."""
     if not (event.content and event.content.parts):
         return None
@@ -316,8 +361,12 @@ def get_auth_config(
     auth_request_function_call: types.FunctionCall,
 ) -> AuthConfig:
     """Extracts the AuthConfig object from the arguments of the auth request function call."""
-    if not auth_request_function_call.args or not (
-        auth_config := auth_request_function_call.args.get('auth_config')
+    if (
+        not auth_request_function_call.args
+        or not isinstance(auth_request_function_call.args, dict)
+        or not (
+            auth_config := auth_request_function_call.args.get('auth_config')
+        )
     ):
         raise ValueError(
             f'Cannot get auth config from function call: {auth_request_function_call}'