[pydantic-ai] 병렬 도구 실행 시 예외 발생 시 형제 태스크 취소 버그 수정

PR 링크: pydantic/pydantic-ai#4502 상태: Merged | 변경: +60 / -1

들어가며

pydantic-ai에서 에이전트가 여러 도구를 병렬로 실행할 때, 한 도구에서 RuntimeError나 ConnectionError 같은 일반 예외가 발생하면 나머지 실행 중인 태스크가 취소되지 않고 "고아(orphaned)" 상태로 남는 치명적인 버그가 있었습니다. 기존 코드는 asyncio.CancelledError만 처리하고 다른 예외는 무시했기 때문입니다. 이 PR은 이 문제를 정확히 해결합니다.

핵심 코드 분석

BaseException 핸들러 추가

Before:

except asyncio.CancelledError as e:
    for task in tasks:
        task.cancel(msg=e.args[0] if len(e.args) != 0 else None)

    raise

기존에는 CancelledError에 대한 처리만 있었고, 그 뒤에 바로 raise가 이어졌습니다. RuntimeError 같은 예외는 이 블록을 건너뛰고 전파되므로, 아직 실행 중인 형제 태스크들이 취소 없이 방치되었습니다.

After:

except asyncio.CancelledError as e:
    for task in tasks:
        task.cancel(msg=e.args[0] if len(e.args) != 0 else None)
    raise
except BaseException:
    # Cancel any still-running sibling tasks so they don't become
    # orphaned asyncio tasks when a non-CancelledError exception
    # propagates out of handle_call_or_result().
    for task in tasks:
        task.cancel()
    raise

except BaseException 블록이 추가되어 모든 종류의 예외에서 형제 태스크를 취소합니다. CancelledError는 별도 핸들러에서 메시지를 전달하며 처리하고, 나머지 모든 예외는 BaseException 핸들러가 잡아서 깔끔하게 정리합니다.

회귀 테스트

async def test_parallel_tool_exception_cancels_sibling_tasks():
    slow_tool_started = asyncio.Event()
    slow_tool_cancelled = asyncio.Event()

    @agent.tool_plain
    async def fast_failing_tool() -> str:
        await asyncio.sleep(0)
        raise RuntimeError('boom')

    @agent.tool_plain
    async def slow_tool() -> str:
        slow_tool_started.set()
        try:
            await asyncio.sleep(10)
        except asyncio.CancelledError:
            slow_tool_cancelled.set()
            raise
        return 'done'

테스트는 빠르게 실패하는 도구와 느린 도구를 병렬 실행하여, RuntimeError 발생 후 느린 도구가 실제로 취소되는지, 고아 태스크가 남지 않는지를 asyncio.all_tasks() 비교로 검증합니다.

왜 이게 좋은가

asyncio에서 고아 태스크는 리소스 누수, 예측 불가능한 부작용, 디버깅이 어려운 간헐적 오류의 원인이 됩니다. 특히 에이전트 프레임워크에서 도구 실행은 외부 API 호출이나 데이터베이스 작업을 포함할 수 있어, 정리되지 않은 태스크가 커넥션 풀 고갈이나 데이터 정합성 문제로 이어질 수 있습니다. 이 수정은 CancelledError와 동일한 패턴으로 BaseException을 처리하여, 예외 종류에 관계없이 항상 깔끔한 정리를 보장합니다.

정리

항목	내용
문제	병렬 도구 실행 시 CancelledError 외 예외에서 형제 태스크 미취소
해결	except BaseException 블록 추가로 모든 예외에서 태스크 취소
영향	고아 태스크로 인한 리소스 누수 및 부작용 방지

참고 자료

• Issue #4423
• Python asyncio Task Cancellation

⚠️ 알림: 이 분석은 AI가 실제 코드 diff를 기반으로 작성했습니다.