refactor: rename transport registration function for clarity and update imports

Author: pi-2r

src/seclab_taskflow_agent/mcp_lifecycle.py

@@ -9,11 +9,11 @@
 
 from __future__ import annotations
 
-__all__ = ["MCP_CLEANUP_TIMEOUT", "build_mcp_servers", "mcp_session_task"]
+__all__ = ["MCP_CLEANUP_TIMEOUT", "build_mcp_servers", "mcp_session_task", "register_transport"]
 
 import asyncio
 import logging
-from types import MappingProxyType
+
 from typing import TYPE_CHECKING, Any, Callable
 
 from agents.mcp import MCPServerSse, MCPServerStdio, MCPServerStreamableHttp, create_static_tool_filter
@@ -48,15 +48,15 @@ def __init__(self, server: MCPNamespaceWrap, process: StreamableMCPThread | None
 MCP_TRANSPORT_REGISTRY: dict[str, MCPServerBuilder] = {}
 
 
-def _register_transport(kind: str) -> Callable[[MCPServerBuilder], MCPServerBuilder]:
+def register_transport(kind: str) -> Callable[[MCPServerBuilder], MCPServerBuilder]:
     """Decorator to register an MCP transport builder."""
     def decorator(builder: MCPServerBuilder) -> MCPServerBuilder:
         MCP_TRANSPORT_REGISTRY[kind] = builder
         return builder
     return decorator
 
 
-@_register_transport("stdio")
+@register_transport("stdio")
 def _build_stdio(tb: str, params: dict[str, Any], tool_filter: Any, client_session_timeout: int, confirms: list[str]) -> MCPServerEntry:
     if params.get("reconnecting", False):
         mcp_server = ReconnectingMCPServerStdio(
@@ -77,7 +77,7 @@ def _build_stdio(tb: str, params: dict[str, Any], tool_filter: Any, client_sessi
     return MCPServerEntry(MCPNamespaceWrap(confirms, mcp_server), None, name=tb)
 
 
-@_register_transport("sse")
+@register_transport("sse")
 def _build_sse(tb: str, params: dict[str, Any], tool_filter: Any, client_session_timeout: int, confirms: list[str]) -> MCPServerEntry:
     mcp_server = MCPServerSse(
         name=tb,
@@ -88,7 +88,7 @@ def _build_sse(tb: str, params: dict[str, Any], tool_filter: Any, client_session
     return MCPServerEntry(MCPNamespaceWrap(confirms, mcp_server), None, name=tb)
 
 
-@_register_transport("streamable")
+@register_transport("streamable")
 def _build_streamable(tb: str, params: dict[str, Any], tool_filter: Any, client_session_timeout: int, confirms: list[str]) -> MCPServerEntry:
     server_proc = None
     if "command" in params:
@@ -115,11 +115,6 @@ def _print_err(line: str) -> None:
     return MCPServerEntry(MCPNamespaceWrap(confirms, mcp_server), server_proc, name=tb)
 
 
-# Freeze the registry after all transports are registered to prevent
-# accidental mutation in tests or at runtime.
-MCP_TRANSPORT_REGISTRY = MappingProxyType(MCP_TRANSPORT_REGISTRY)
-
-
 def build_mcp_servers(
     available_tools: AvailableTools,
     toolboxes: list[str],

src/seclab_taskflow_agent/runner.py

@@ -33,7 +33,7 @@
 from agents.extensions.handoff_prompt import prompt_with_handoff_instructions
 from openai import APIConnectionError, APITimeoutError, BadRequestError, RateLimitError
 from openai.types.responses import ResponseTextDeltaEvent
-from tenacity import AsyncRetrying, retry, retry_if_exception_type, stop_after_attempt, wait_exponential
+from tenacity import AsyncRetrying, retry, retry_if_exception_type, retry_if_not_exception_type, stop_after_attempt, wait_exponential
 
 from .agent import TaskAgent, TaskAgentHooks, TaskRunHooks
 from .capi import get_default_model