Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions .github/actions/spelling/allow.txt
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
a2a

Check warning on line 1 in .github/actions/spelling/allow.txt

View workflow job for this annotation

GitHub Actions / Check Spelling

Ignoring entry because it contains non-alpha characters (non-alpha-in-dictionary)
A2A

Check warning on line 2 in .github/actions/spelling/allow.txt

View workflow job for this annotation

GitHub Actions / Check Spelling

Ignoring entry because it contains non-alpha characters (non-alpha-in-dictionary)
A2AFastAPI

Check warning on line 3 in .github/actions/spelling/allow.txt

View workflow job for this annotation

GitHub Actions / Check Spelling

Ignoring entry because it contains non-alpha characters (non-alpha-in-dictionary)
AAgent
Expand Down Expand Up @@ -28,7 +28,7 @@
AUser
autouse
backticks
base64url

Check warning on line 31 in .github/actions/spelling/allow.txt

View workflow job for this annotation

GitHub Actions / Check Spelling

Ignoring entry because it contains non-alpha characters (non-alpha-in-dictionary)
buf
bufbuild
cla
Expand All @@ -43,7 +43,7 @@
drivername
DSNs
dunders
ES256

Check warning on line 46 in .github/actions/spelling/allow.txt

View workflow job for this annotation

GitHub Actions / Check Spelling

Ignoring entry because it contains non-alpha characters (non-alpha-in-dictionary)
euo
EUR
evt
Expand All @@ -58,8 +58,8 @@
gle
GVsb
hazmat
HS256

Check warning on line 61 in .github/actions/spelling/allow.txt

View workflow job for this annotation

GitHub Actions / Check Spelling

Ignoring entry because it contains non-alpha characters (non-alpha-in-dictionary)
HS384

Check warning on line 62 in .github/actions/spelling/allow.txt

View workflow job for this annotation

GitHub Actions / Check Spelling

Ignoring entry because it contains non-alpha characters (non-alpha-in-dictionary)
ietf
importlib
initdb
Expand Down Expand Up @@ -100,10 +100,10 @@
Oneof
OpenAPI
openapiv
openapiv2

Check warning on line 103 in .github/actions/spelling/allow.txt

View workflow job for this annotation

GitHub Actions / Check Spelling

Ignoring entry because it contains non-alpha characters (non-alpha-in-dictionary)
opensource
otherurl
pb2

Check warning on line 106 in .github/actions/spelling/allow.txt

View workflow job for this annotation

GitHub Actions / Check Spelling

Ignoring entry because it contains non-alpha characters (non-alpha-in-dictionary)
podman
Podman
poolclass
Expand All @@ -126,7 +126,7 @@
respx
resub
rmi
RS256

Check warning on line 129 in .github/actions/spelling/allow.txt

View workflow job for this annotation

GitHub Actions / Check Spelling

Ignoring entry because it contains non-alpha characters (non-alpha-in-dictionary)
RUF
SECP256R1
SFIXED
Expand All @@ -150,4 +150,5 @@
typeerror
UIDs
vulnz
GC
whl
85 changes: 52 additions & 33 deletions src/a2a/server/agent_execution/active_task.py
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,7 @@
from __future__ import annotations

import asyncio
import contextlib
import logging
import uuid

Expand Down Expand Up @@ -577,6 +578,18 @@ async def _run_consumer(self) -> None:
async with self._lock:
self._reference_count -= 1
logger.debug('Consumer[%s]: Finishing', self._task_id)

# Cancel and await the producer before cleanup to prevent
# asyncio "Task was destroyed but it is pending!" warnings.
# Without this, the producer could still be parked on
# _request_queue.get() or inside its finally block when
# _maybe_cleanup() releases the ActiveTask reference,
# causing the still-pending task to be garbage-collected.
if self._producer_task and not self._producer_task.done():
self._producer_task.cancel()
with contextlib.suppress(asyncio.CancelledError):
await self._producer_task

await self._maybe_cleanup()

async def subscribe(
Expand Down Expand Up @@ -688,43 +701,49 @@ async def cancel(self, call_context: ServerCallContext) -> Task:
"""
logger.debug('Cancel[%s]: Cancelling task', self._task_id)

# TODO: Conflicts with call_context on the pending request.
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,
)
# Save and restore _call_context to avoid races with the producer
# loop which also sets _call_context for incoming requests.
# See the TODO at the producer loop line ~512.
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
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
)
except Exception as e:
logger.exception(
'Cancel[%s]: Agent cancel failed', 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,
)
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,
)

await self._is_finished.wait()
task = await self._task_manager.get_task()
await self._is_finished.wait()
task = await self._task_manager.get_task()
finally:
self._task_manager._call_context = saved_call_context
Comment on lines +707 to +746

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

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

if not task:
raise RuntimeError('Task should have been created')
return task
Expand Down
15 changes: 14 additions & 1 deletion src/a2a/server/tasks/base_push_notification_sender.py
Original file line number Diff line number Diff line change
Expand Up @@ -82,10 +82,23 @@ async def _dispatch_notification(
) -> bool:
url = push_info.url
try:
headers = None
headers: dict[str, str] | None = None
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:
if headers is None:
headers = {}
headers['Authorization'] = f'{scheme} {credentials}'
Comment on lines +85 to +100

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

low

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}'


response = await self._client.post(
url,
json=MessageToDict(to_stream_response(event)),
Expand Down
154 changes: 154 additions & 0 deletions tests/server/agent_execution/test_active_task.py
Original file line number Diff line number Diff line change
Expand Up @@ -895,3 +895,157 @@ async def execute_mock(req, q):
assert len(events) == 0

await active_task.cancel(request_context)


@pytest.mark.timeout(5)
@pytest.mark.asyncio
async def test_producer_cancelled_on_normal_completion():
"""Verify producer is cancelled and awaited before cleanup to prevent GC warnings.

Regression test for #1121: when the consumer finishes first (normal
completion path), the producer must be cancelled and awaited before
_maybe_cleanup() releases the ActiveTask reference, otherwise asyncio
logs "Task was destroyed but it is pending!".
"""
agent_executor = Mock()
task_manager = Mock()
cleanup_called = False

def on_cleanup(_task: ActiveTask) -> None:
nonlocal cleanup_called
cleanup_called = True

active_task = ActiveTask(
agent_executor=agent_executor,
task_id='test-task-id',
task_manager=task_manager,
on_cleanup=on_cleanup,
)

# Simulate a producer that blocks waiting for a request.
# The consumer will finish first (after processing all events),
# triggering normal teardown.
execute_started = asyncio.Event()
execute_barrier = asyncio.Event()

async def execute_mock(req, q):
execute_started.set()
# Block until released — simulates producer waiting for next request
await execute_barrier.wait()

agent_executor.execute = AsyncMock(side_effect=execute_mock)
agent_executor.cancel = AsyncMock()
task_manager.get_task = AsyncMock(
return_value=Task(
id='test-task-id',
status=TaskStatus(state=TaskState.TASK_STATE_WORKING),
)
)
task_manager.save_task_event = AsyncMock()
task_manager.ensure_task_id = AsyncMock(
return_value=Task(
id='test-task-id',
status=TaskStatus(state=TaskState.TASK_STATE_WORKING),
)
)
task_manager.process = AsyncMock(side_effect=lambda x: x)

request_context = Mock(spec=RequestContext)
request_context.call_context = ServerCallContext()
request_context.context_id = 'test-context'
request_context.message = None

await active_task.enqueue_request(request_context)
await active_task.start(
call_context=ServerCallContext(), create_task_if_missing=True
)

# Wait for producer to be executing
await execute_started.wait()

# Cancel the consumer to force consumer teardown while producer is pending.
# This simulates the normal-completion race: consumer finishes first,
# triggering _maybe_cleanup while producer is still parked.
if active_task._consumer_task:
active_task._consumer_task.cancel()
try:
await active_task._consumer_task
except asyncio.CancelledError:
pass

# Release the producer so it can complete its CancelledError/finally path
execute_barrier.set()

# Wait for producer to finish
await active_task._is_finished.wait()

# Verify producer was properly completed, not left pending
assert active_task._producer_task is not None
assert active_task._producer_task.done(), (
'Producer task should be done after consumer teardown'
)
assert cleanup_called, (
'on_cleanup should be called after producer is drained'
)


@pytest.mark.timeout(5)
@pytest.mark.asyncio
async def test_cancel_restores_call_context():
"""Verify cancel() restores _call_context after completion.

Regression test: cancel() temporarily sets _call_context to its
own call_context for task operations, then restores the original
value to prevent races with the producer loop.
"""
agent_executor = Mock()
task_manager = Mock()
task_manager.get_task = AsyncMock()
original_context = ServerCallContext()
task_manager._call_context = original_context

active_task = ActiveTask(
agent_executor=agent_executor,
task_id='test-task-id',
task_manager=task_manager,
)

stop_event = asyncio.Event()

async def execute_mock(req, q):
await stop_event.wait()

agent_executor.execute = AsyncMock(side_effect=execute_mock)
agent_executor.cancel = AsyncMock()
task_manager.get_task.side_effect = [
Task(
id='test-task-id',
status=TaskStatus(state=TaskState.TASK_STATE_WORKING),
)
] * 20
task_manager.save_task_event = AsyncMock()
task_manager.ensure_task_id = AsyncMock(
return_value=Task(
id='test-task-id',
status=TaskStatus(state=TaskState.TASK_STATE_WORKING),
)
)
task_manager.process = AsyncMock(side_effect=lambda x: x)

request_context = Mock(spec=RequestContext)
request_context.call_context = original_context
request_context.context_id = 'test-context'
request_context.message = None

await active_task.enqueue_request(request_context)
await active_task.start(
call_context=original_context, create_task_if_missing=True
)
await asyncio.sleep(0.1)

cancel_context = ServerCallContext()
await active_task.cancel(cancel_context)
stop_event.set()

# Verify call_context was restored
assert task_manager._call_context is original_context
Loading
Loading