본문으로 건너뛰기

[pydantic-ai] Temporal/DBOS MCP 서버에서 매번 도구 목록을 다시 가져오는 문제 수정

PR 링크: pydantic/pydantic-ai#4331 상태: Merged | 변경: +3008 / -45

들어가며

pydantic-ai의 durable execution (Temporal, DBOS) 환경에서 MCP 서버와 통신할 때, 매 activity/step마다 도구 목록을 새로 가져오는 비효율이 있었습니다. MCP 서버의 cache_tools 설정이 존재했지만 래퍼 계층에서 이를 활용하지 못했습니다. 이 PR은 래퍼 수준에서 도구 정의 캐시를 추가하여 불필요한 MCP 연결을 제거합니다.

핵심 코드 분석

DBOSMCPServer 캐시 추가

Before:

class DBOSMCPServer(DBOSMCPToolset[AgentDepsT]):
    def __init__(self, ...):
        super().__init__(...)

    def tool_for_tool_def(self, tool_def: ToolDefinition) -> ToolsetTool[AgentDepsT]:
        assert isinstance(self.wrapped, MCPServer)
        return self.wrapped.tool_for_tool_def(tool_def)

매 get_tools 호출 시 부모 클래스의 DBOS step을 통해 MCP 서버에 접속하여 도구 목록을 가져왔습니다.

After:

class DBOSMCPServer(DBOSMCPToolset[AgentDepsT]):
    def __init__(self, ...):
        super().__init__(...)
        self._cached_tool_defs: dict[str, ToolDefinition] | None = None

    async def get_tools(self, ctx: RunContext[AgentDepsT]) -> dict[str, ToolsetTool[AgentDepsT]]:
        if self._server.cache_tools and self._cached_tool_defs is not None:
            return {name: self.tool_for_tool_def(td) for name, td in self._cached_tool_defs.items()}

        result = await super().get_tools(ctx)
        if self._server.cache_tools:
            self._cached_tool_defs = {name: tool.tool_def for name, tool in result.items()}
        return result

cache_tools가 활성화되어 있으면 첫 호출 결과를 _cached_tool_defs에 저장하고, 이후 호출에서는 캐시된 ToolDefinition으로 ToolsetTool을 구성합니다. 동일한 패턴이 TemporalMCPServer에도 적용되었습니다.

MCPServer의 list_tools 최적화

# Before: 캐시 확인이 `async with self:` 내부에 있었음
async with self:
    if self.cache_tools:
        if self._cached_tools is not None:
            return self._cached_tools
        result = await self._client.list_tools()

# After: 캐시 확인을 컨텍스트 진입 전에 수행
if self.cache_tools and self._cached_tools is not None:
    return self._cached_tools
async with self:
    result = await self._client.list_tools()
    if self.cache_tools:
        self._cached_tools = result.tools
    return result.tools

캐시 히트 시 MCP 연결 컨텍스트(async with self)에 진입하지 않아 더 효율적입니다.

왜 이게 좋은가

Durable execution 환경에서 각 activity/step은 독립적인 실행 단위입니다. 도구 목록이 변경되지 않는 한 매번 MCP 서버에 연결하여 목록을 가져오는 것은 순수한 낭비입니다. 특히 네트워크 왕복(round-trip)이 필요한 원격 MCP 서버의 경우, 이 캐시가 워크플로우 전체 지연 시간을 크게 줄일 수 있습니다. cache_tools=False 설정으로 동적 도구 변경이 필요한 경우를 지원하는 안전장치도 마련되어 있습니다.

정리

항목 내용
문제 매 activity/step마다 MCP 서버에서 도구 목록 재조회
해결 래퍼 수준 _cached_tool_defs 캐시 추가
적용 DBOSMCPServer, TemporalMCPServer 동일 패턴

참고 자료

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

댓글

관련 포스트

PR Analysis 의 다른글