[pydantic-ai] RunUsage.tool_calls 병렬 실행 시 과소 집계 버그 수정 (asyncio.Lock)

PR 링크: pydantic/pydantic-ai#3133 상태: Merged (이후 #3174에서 revert됨) | 변경: +48 / -4

들어가며

Pydantic AI가 LLM으로부터 여러 tool call을 동시에 받으면, asyncio.create_task()로 병렬 실행합니다. 각 task가 완료 후 usage.tool_calls += 1을 수행하는데, asyncio에서 await 지점 사이의 read-modify-write가 인터리빙되면 카운터가 과소 집계될 수 있었습니다.

핵심 코드 분석

asyncio.Lock으로 카운터 보호

Before (_tool_manager.py):

tool_result = await self._call_tool(call, allow_partial, wrap_validation_errors)
usage.tool_calls += 1

After:

@cached_property
def _usage_lock(self) -> asyncio.Lock:
    """Lock to prevent race conditions when incrementing usage.tool_calls."""
    return asyncio.Lock()

# ...
tool_result = await self._call_tool(call, allow_partial, wrap_validation_errors)
async with self._usage_lock:
    usage.tool_calls += 1

cached_property로 Lock을 지연 생성하여, Lock이 필요하지 않은 순차 실행 경로에서는 오버헤드가 없습니다.

재현 테스트

async def test_race_condition_parallel_tool_calls():
    for iteration in range(20):  # 반복하여 race condition 재현 확률 증가
        # 10개의 병렬 tool call 생성
        return ModelResponse(parts=[ToolCallPart('tool_a', {}, f'call_{i}') for i in range(10)])
        
        @agent.tool_plain
        async def tool_a() -> str:
            await asyncio.sleep(0.0001)  # task 인터리빙 유도
            await asyncio.sleep(0.0001)
            return 'result'
        
        actual = result.usage().tool_calls
        assert actual == 10  # Lock 없이는 간헐적으로 실패

왜 이게 좋은가

• 10개의 병렬 tool call에서 실제 실행 횟수가 정확히 기록됩니다.
• 비용 추적, rate limiting, usage limit 검사 등이 정확한 tool_calls 값에 의존하므로, 이 수정은 실질적 영향이 있습니다.

참고: 이 PR은 이후 #3174에서 revert되었고, #2978에서 더 근본적인 접근(tool call 실행 전 사전 검증)으로 대체되었습니다.

정리

• asyncio에서 +=는 원자적이지 않다: await 지점에서 task가 전환되므로, 공유 변수의 read-modify-write는 race condition을 유발합니다.
• 간헐적 버그는 반복 테스트로 잡아라: 20회 반복과 의도적인 asyncio.sleep()으로 race condition 재현 확률을 높인 테스트 전략이 효과적입니다.

참고 자료

• pydantic/pydantic-ai#3133 — PR 전체 diff
• pydantic/pydantic-ai#3174 — Revert PR
• pydantic/pydantic-ai#2978 — 최종 수정

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