fix(server): save and restore call_context in cancel() to prevent race with producer

[isheng-eqi]

Summary

Fix the _call_context race condition in ActiveTask.cancel() — use a save/restore pattern so the cancel path doesn't overwrite the shared _call_context while the producer loop is using it.

Problem

cancel() directly set self._task_manager._call_context = call_context, which races with the producer loop at line ~512 that also sets _call_context for each incoming request:

# cancel() line 705:
self._task_manager._call_context = call_context  

# _run_producer loop line 512:
self._task_manager._call_context = request_context.call_context

_call_context is used by task_store.get() and task_store.save() for authorization checks. When cancel and a normal request happen concurrently, _call_context gets overwritten — the wrong caller identity could be used for task store operations. A TODO at the site documented this known issue.

Fix

Wrap cancel()'s body in a save/restore pattern:

saved_call_context = self._task_manager._call_context
try:
    self._task_manager._call_context = call_context
    # ... all cancel operations ...
finally:
    self._task_manager._call_context = saved_call_context

This ensures the cancel uses its own call_context for its operations, then restores whatever was there before — whether that's the producer's or the start() original.

Test Plan

• test_cancel_restores_call_context: creates ActiveTask with original call_context, calls cancel() with different context, verifies original is restored
• ruff check + ruff format + ty check pass

[gemini-code-assist]

Code Review

This pull request addresses asyncio warnings by ensuring the producer task is cancelled and awaited before cleanup, and introduces a save/restore pattern for the call context during task cancellation to prevent races. It also adds support for push notification authentication by appending an Authorization header when configured, along with corresponding unit tests. The review feedback recommends narrowing the window of the call context save/restore pattern to avoid concurrency issues during long-running awaits, and simplifying the header initialization logic by using an empty dictionary instead of checking for None.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

Modifying the shared self._task_manager._call_context with a save/restore pattern across a long-running await self._is_finished.wait() is highly vulnerable to race conditions. If multiple cancel() calls or other concurrent operations occur, the shared context can be corrupted or restored to the wrong value.

To minimize the race window and ensure robustness, we should only set _call_context temporarily during the active cancellation phase, restore it before awaiting _is_finished.wait(), and then temporarily set it again to retrieve the final task state.

        saved_call_context = self._task_manager._call_context
        try:
            self._task_manager._call_context = call_context

            task = await self._task_manager.get_task()
            request_context = RequestContext(
                call_context=call_context,
                task_id=self._task_id,
                context_id=task.context_id if task else None,
                task=task,
            )

            async with self._lock:
                if not self._is_finished.is_set() and self._producer_task:
                    logger.debug(
                        'Cancel[%s]: Cancelling producer task', self._task_id
                    )
                    self._producer_task.cancel()
                    try:
                        await self._agent_executor.cancel(
                            request_context, self._event_queue_agent
                        )
                    except Exception as e:
                        logger.exception(
                            'Cancel[%s]: Agent cancel failed', self._task_id
                        )
                        await self._mark_task_as_failed(e)
                        raise
                else:
                    logger.debug(
                        'Cancel[%s]: Task already finished [%s] or producer not started [%s], not cancelling',
                        self._task_id,
                        self._is_finished.is_set(),
                        self._producer_task,
                    )
        finally:
            self._task_manager._call_context = saved_call_context

        await self._is_finished.wait()

        saved_call_context = self._task_manager._call_context
        try:
            self._task_manager._call_context = call_context
            task = await self._task_manager.get_task()
        finally:
            self._task_manager._call_context = saved_call_context

[gemini-code-assist]

Instead of initializing headers as None and performing repetitive None checks, we can initialize headers as an empty dictionary and conditionally populate it. This simplifies the logic and improves readability.

            headers: dict[str, str] = {}
            if push_info.token:
                headers['X-A2A-Notification-Token'] = push_info.token

            # Add Authorization header when authentication is configured.
            # The A2A spec requires servers to authenticate push notifications
            # when the client provides an authentication scheme in
            # PushNotificationConfig.
            if push_info.HasField('authentication'):
                auth = push_info.authentication
                scheme = auth.scheme
                credentials = auth.credentials
                if scheme and credentials:
                    headers['Authorization'] = f'{scheme} {credentials}'