fix: store calendar OAuth credentials after authentication

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

Author: anxkhn

samples/python/agents/birthday_planner_adk/calendar_agent/adk_agent_executor.py

@@ -89,16 +89,10 @@ async def _process_request(
             # 2. The function call required authorization.
             # Ideally we'd have a way to interpret whether the response is a completion for the
             # task or requires follow-up, but I'm not going to bother just yet.
-            if auth_request_function_call := get_auth_request_function_call(
-                event
-            ):
+            if auth_request_function_call := get_auth_request_function_call(event):
                 # Gather details, then suspend.
-                auth_details = self._prepare_auth_request(
-                    auth_request_function_call
-                )
-                logger.debug(
-                    'Yielding auth required response: %s', auth_details.uri
-                )
+                auth_details = self._prepare_auth_request(auth_request_function_call)
+                logger.debug('Yielding auth required response: %s', auth_details.uri)
                 await task_updater.update_status(
                     TaskState.auth_required,
                     message=new_agent_text_message(
@@ -127,9 +121,7 @@ async def _process_request(
 
         if auth_details:
             # After auth is received, we can continue processing this request.
-            await self._complete_auth_processing(
-                context, auth_details, task_updater
-            )
+            await self._complete_auth_processing(context, auth_details, task_updater)
 
     def _prepare_auth_request(
         self, auth_request_function_call: types.FunctionCall
@@ -148,9 +140,7 @@ def _prepare_auth_request(
         oauth2_config = auth_config.exchanged_auth_credential.oauth2
         base_auth_uri = oauth2_config.auth_uri
         if not base_auth_uri:
-            raise ValueError(
-                f'Cannot get auth uri from auth config: {auth_config}'
-            )
+            raise ValueError(f'Cannot get auth uri from auth config: {auth_config}')
         redirect_uri = f'{self._card.url}authenticate'
         oauth2_config.redirect_uri = redirect_uri
         state_token = oauth2_config.state
@@ -194,9 +184,7 @@ async def _complete_auth_processing(
             ),
         )
         del self._awaiting_auth[auth_details.state]
-        oauth2_config = (
-            auth_details.auth_config.exchanged_auth_credential.oauth2
-        )
+        oauth2_config = auth_details.auth_config.exchanged_auth_credential.oauth2
         oauth2_config.auth_response_uri = auth_uri
         auth_content = types.UserContent(
             parts=[
@@ -210,20 +198,20 @@ async def _complete_auth_processing(
             ]
         )
         await self._process_request(auth_content, context, task_updater)
-        # Extract the stored credential.
-        if context.call_context and context.call_context.user.is_authenticated:
-            await self._store_user_auth(
-                context,
-                auth_details.auth_config.auth_scheme,
-                auth_details.auth_config.raw_auth_credential,
-            )
+        # Always hoist the session credential. The documented OAuth redirect
+        # has no JWT, so call_context stays unauthenticated.
+        await self._store_user_auth(
+            context,
+            auth_details.auth_config.auth_scheme,
+            auth_details.auth_config.raw_auth_credential,
+        )
 
     async def execute(
         self,
         context: RequestContext,
         event_queue: EventQueue,
-    ):
-        # Run the agent until either complete or the task is suspended.
+    ) -> None:
+        """Run the agent until the task completes or is suspended."""
         updater = TaskUpdater(event_queue, context.task_id, context.context_id)
         # Immediately notify that the task is submitted.
         if not context.current_task:
@@ -238,11 +226,12 @@ async def execute(
         )
         logger.debug('[Calendar] execute exiting')
 
-    async def cancel(self, context: RequestContext, event_queue: EventQueue):
-        # Ideally: kill any ongoing tasks.
+    async def cancel(self, context: RequestContext, event_queue: EventQueue) -> None:
+        """Cancel is not supported for the calendar agent."""
         raise ServerError(error=UnsupportedOperationError())
 
-    async def on_auth_callback(self, state: str, uri: str):
+    async def on_auth_callback(self, state: str, uri: str) -> None:
+        """Resume an in-flight OAuth callback."""
         self._awaiting_auth[state].set_result(uri)
 
     async def _upsert_session(self, context: RequestContext) -> Session:
@@ -262,9 +251,9 @@ async def _upsert_session(self, context: RequestContext) -> Session:
         return await self._ensure_auth(session)
 
     async def _ensure_auth(self, session: Session) -> Session:
-        if (
-            stored_cred := self._credentials.get(session.user_id)
-        ) and not session.state.get(stored_cred.key):
+        if (stored_cred := self._credentials.get(session.user_id)) and not session.state.get(
+            stored_cred.key
+        ):
             event_action = EventActions(
                 state_delta={
                     stored_cred.key: stored_cred.credential,
@@ -301,10 +290,8 @@ async def _store_user_auth(
         )
         stored_credential = session.state.get(credential_key)
         if stored_credential:
-            self._credentials[context.call_context.user.user_name] = (
-                StoredCredential(
-                    key=credential_key, credential=stored_credential
-                )
+            self._credentials[session.user_id] = StoredCredential(
+                key=credential_key, credential=stored_credential
             )
 
 
@@ -321,15 +308,11 @@ def convert_a2a_part_to_genai(part: Part) -> types.Part:
     if isinstance(part, FilePart):
         if isinstance(part.file, FileWithUri):
             return types.Part(
-                file_data=types.FileData(
-                    file_uri=part.file.uri, mime_type=part.file.mime_type
-                )
+                file_data=types.FileData(file_uri=part.file.uri, mime_type=part.file.mime_type)
             )
         if isinstance(part.file, FileWithBytes):
             return types.Part(
-                inline_data=types.Blob(
-                    data=part.file.bytes, mime_type=part.file.mime_type
-                )
+                inline_data=types.Blob(data=part.file.bytes, mime_type=part.file.mime_type)
             )
         raise ValueError(f'Unsupported file type: {type(part.file)}')
     raise ValueError(f'Unsupported part type: {type(part)}')
@@ -390,7 +373,5 @@ def get_auth_config(
     if not auth_request_function_call.args or not (
         auth_config := auth_request_function_call.args.get('authConfig')
     ):
-        raise ValueError(
-            f'Cannot get auth config from function call: {auth_request_function_call}'
-        )
+        raise ValueError(f'Cannot get auth config from function call: {auth_request_function_call}')
     return AuthConfig.model_validate(auth_config)

samples/python/agents/birthday_planner_adk/calendar_agent/tests/__init__.py

@@ -0,0 +1,158 @@
+import asyncio
+import importlib.util
+import sys
+
+from pathlib import Path
+from types import ModuleType
+from unittest.mock import AsyncMock, MagicMock, patch
+
+
+_AGENT_DIR = Path(__file__).resolve().parents[1]
+if str(_AGENT_DIR) not in sys.path:
+    sys.path.insert(0, str(_AGENT_DIR))
+
+
+class _StubAgentExecutor:
+    """Stand-in for a2a.server.agent_execution.AgentExecutor."""
+
+
+def _ensure_module(name: str) -> ModuleType:
+    module = sys.modules.get(name)
+    if module is not None:
+        return module
+    module = ModuleType(name)
+    module.__path__ = []
+    sys.modules[name] = module
+    parent_name, _, attr = name.rpartition('.')
+    if parent_name:
+        setattr(_ensure_module(parent_name), attr, module)
+    return module
+
+
+def _missing(name: str) -> bool:
+    try:
+        return importlib.util.find_spec(name) is None
+    except ModuleNotFoundError:
+        return True
+
+
+def _set_attr(module: ModuleType, name: str, value: object) -> None:
+    if not hasattr(module, name):
+        setattr(module, name, value)
+
+
+def _install_import_stubs() -> None:
+    """Let this test import the executor without the sample's runtime extras."""
+    agent_execution = _ensure_module('a2a.server.agent_execution')
+    _set_attr(agent_execution, 'AgentExecutor', _StubAgentExecutor)
+    _set_attr(agent_execution, 'RequestContext', MagicMock)
+    _set_attr(
+        _ensure_module('a2a.server.agent_execution.context'),
+        'RequestContext',
+        agent_execution.RequestContext,
+    )
+    _set_attr(_ensure_module('a2a.server.events.event_queue'), 'EventQueue', MagicMock)
+    _set_attr(_ensure_module('a2a.server.tasks'), 'TaskUpdater', MagicMock)
+    a2a_types = _ensure_module('a2a.types')
+    for type_name in (
+        'AgentCard',
+        'FilePart',
+        'FileWithBytes',
+        'FileWithUri',
+        'Part',
+        'TextPart',
+        'UnsupportedOperationError',
+    ):
+        _set_attr(a2a_types, type_name, MagicMock)
+
+    class _TaskState:
+        working = 'working'
+        failed = 'failed'
+        auth_required = 'auth_required'
+
+    a2a_types.TaskState = _TaskState
+    _set_attr(_ensure_module('a2a.utils.errors'), 'ServerError', Exception)
+    _set_attr(
+        _ensure_module('a2a.utils.message'),
+        'new_agent_text_message',
+        MagicMock(),
+    )
+    if _missing('google.adk'):
+        _ensure_module('google.adk').Runner = MagicMock
+        adk_auth = _ensure_module('google.adk.auth')
+        for type_name in ('AuthConfig', 'AuthCredential', 'AuthScheme'):
+            setattr(adk_auth, type_name, MagicMock)
+        adk_events = _ensure_module('google.adk.events')
+        for type_name in ('Event', 'EventActions'):
+            setattr(adk_events, type_name, MagicMock)
+        _ensure_module('google.adk.sessions').Session = MagicMock
+        _ensure_module(
+            'google.adk.tools.openapi_tool.openapi_spec_parser.tool_auth_handler'
+        ).ToolContextCredentialStore = MagicMock
+    if _missing('google.genai'):
+        _ensure_module('google.genai').types = MagicMock()
+
+
+_install_import_stubs()
+
+from adk_agent_executor import ADKAgentExecutor, ADKAuthDetails
+
+
+def test_complete_auth_processing_stores_credential_when_unauthenticated() -> None:
+    """OAuth callback without a JWT must still cache the session credential."""
+    asyncio.run(_run_complete_auth_processing_unauthenticated())
+
+
+async def _run_complete_auth_processing_unauthenticated() -> None:
+    runner = MagicMock()
+    runner.app_name = 'Calendar Agent'
+    executor = ADKAgentExecutor(runner, MagicMock())
+
+    credential = object()
+    session = MagicMock()
+    session.user_id = 'anonymous'
+    session.state = {'calendar-oauth-key': credential}
+
+    user = MagicMock()
+    user.is_authenticated = False
+    user.user_name = 'jwt-user'
+    call_context = MagicMock()
+    call_context.user = user
+    context = MagicMock()
+    context.call_context = call_context
+    context.context_id = 'ctx-unauthenticated'
+
+    future = asyncio.get_running_loop().create_future()
+    future.set_result('http://localhost:10007/authenticate?code=abc&state=xyz')
+
+    auth_config = MagicMock()
+    auth_config.exchanged_auth_credential.oauth2 = MagicMock()
+    auth_config.auth_scheme = MagicMock()
+    auth_config.raw_auth_credential = MagicMock()
+    auth_details = ADKAuthDetails(
+        state='xyz',
+        uri='https://accounts.google.com/o/oauth2/auth',
+        future=future,
+        auth_config=auth_config,
+        auth_request_function_call_id='fn-1',
+    )
+    executor._awaiting_auth[auth_details.state] = future  # noqa: SLF001
+
+    task_updater = MagicMock()
+    task_updater.update_status = AsyncMock()
+
+    with (
+        patch.object(executor, '_process_request', new_callable=AsyncMock),
+        patch.object(executor, '_upsert_session', new_callable=AsyncMock, return_value=session),
+        patch('adk_agent_executor.ToolContextCredentialStore') as store_cls,
+    ):
+        store_cls.return_value.get_credential_key.return_value = 'calendar-oauth-key'
+        await executor._complete_auth_processing(  # noqa: SLF001
+            context, auth_details, task_updater
+        )
+
+    assert session.user_id in executor._credentials  # noqa: S101, SLF001
+    stored = executor._credentials[session.user_id]  # noqa: SLF001
+    assert stored.key == 'calendar-oauth-key'  # noqa: S101
+    assert stored.credential is credential  # noqa: S101
+    assert user.user_name not in executor._credentials  # noqa: S101, SLF001