"""
Wait idle workers tool.
Provides wait_idle_workers for blocking until workers become idle.
"""
from typing import TYPE_CHECKING
from mcp.server.fastmcp import Context, FastMCP
from mcp.server.session import ServerSession
if TYPE_CHECKING:
from ..server import AppContext
from ..idle_detection import (
wait_for_all_idle as wait_for_all_idle_impl,
wait_for_any_idle as wait_for_any_idle_impl,
SessionInfo,
)
from ..registry import SessionStatus
from ..utils import error_response, HINTS
def register_tools(mcp: FastMCP) -> None:
"""Register wait_idle_workers tool on the MCP server."""
@mcp.tool()
async def wait_idle_workers(
ctx: Context[ServerSession, "AppContext"],
session_ids: list[str],
mode: str | None = "all",
timeout: float | None = 600.0,
poll_interval: float | None = 2.0,
) -> dict:
"""
Wait for worker sessions to become idle.
Unified tool for waiting on one or more workers. Supports two modes:
- "all": Wait until ALL workers are idle (default, for fan-out/fan-in)
- "any": Return as soon as ANY worker becomes idle (for pipelines)
Args:
session_ids: List of session IDs to wait on (accepts 1 or more).
Accepts internal IDs, terminal IDs, or worker names.
mode: "all" or "any" - default "all"
timeout: Maximum seconds to wait (default 10 minutes)
poll_interval: Seconds between checks (default 2)
Returns:
Dict with:
- session_ids: The session IDs that were requested
- idle_session_ids: List of sessions that are idle
- all_idle: Whether all sessions are idle
- waiting_on: Sessions still working (if timed out)
- mode: The mode used
- waited_seconds: How long we waited
- timed_out: Whether we hit the timeout
"""
# Handle None values from MCP clients that send explicit null for omitted params
mode = mode or "all"
timeout = timeout if timeout is not None else 600.0
poll_interval = poll_interval if poll_interval is not None else 2.0
app_ctx = ctx.request_context.lifespan_context
registry = app_ctx.registry
# Validate inputs
if not session_ids:
return error_response(
"session_ids is required and must contain at least one session ID",
hint=HINTS["registry_empty"],
)
# Validate mode
if mode not in ("all", "any"):
return error_response(
f"Invalid mode: {mode}. Must be 'all' or 'any'",
)
# Look up sessions and build SessionInfo list
# Uses resolve() to accept internal ID, terminal ID, or name
session_infos = []
missing_sessions = []
missing_jsonl = []
for session_id in session_ids:
session = registry.resolve(session_id)
if not session:
missing_sessions.append(session_id)
continue
jsonl_path = session.get_jsonl_path()
if not jsonl_path:
missing_jsonl.append(session_id)
continue
session_infos.append(SessionInfo(
jsonl_path=jsonl_path,
session_id=session.session_id, # Must use internal ID to match stop hook marker
agent_type=session.agent_type,
))
# Report any missing sessions/files
if missing_sessions:
return error_response(
f"Sessions not found: {', '.join(missing_sessions)}",
hint=HINTS["session_not_found"],
)
if missing_jsonl:
return error_response(
f"No JSONL files for: {', '.join(missing_jsonl)}",
hint=HINTS["no_jsonl_file"],
)
# Wait based on mode
if mode == "any":
result = await wait_for_any_idle_impl(
sessions=session_infos,
timeout=timeout,
poll_interval=poll_interval,
)
# Convert to common format
idle_session_ids = [result["idle_session_id"]] if result["idle_session_id"] else []
return {
"session_ids": session_ids,
"idle_session_ids": idle_session_ids,
"all_idle": len(idle_session_ids) == len(session_ids),
"waiting_on": [s for s in session_ids if s not in idle_session_ids],
"mode": mode,
"waited_seconds": result["waited_seconds"],
"timed_out": result["timed_out"],
}
else:
# mode == "all"
result = await wait_for_all_idle_impl(
sessions=session_infos,
timeout=timeout,
poll_interval=poll_interval,
)
# Update statuses for idle sessions
for session_id in result["idle_session_ids"]:
registry.update_status(session_id, SessionStatus.READY)
return {
"session_ids": session_ids,
"idle_session_ids": result["idle_session_ids"],
"all_idle": result["all_idle"],
"waiting_on": result["waiting_on"],
"mode": mode,
"waited_seconds": result["waited_seconds"],
"timed_out": result["timed_out"],
}
