"""
Close workers tool.
Provides close_workers for gracefully terminating Claude Code worker sessions.
"""
import asyncio
import logging
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 ..iterm_utils import CODEX_PRE_ENTER_DELAY
from ..registry import SessionRegistry, SessionStatus
from ..worktree import WorktreeError, remove_worktree
from ..utils import error_response, HINTS
logger = logging.getLogger("maniple")
def _compute_prompt_delay(text: str, agent_type: str) -> float:
"""Compute a safe delay before sending Enter for a prompt."""
line_count = text.count("\n")
char_count = len(text)
if line_count > 0:
paste_delay = min(2.0, 0.1 + (line_count * 0.01) + (char_count / 1000 * 0.05))
else:
paste_delay = 0.05
if agent_type == "codex":
return max(CODEX_PRE_ENTER_DELAY, paste_delay)
return paste_delay
async def _send_prompt(backend, session, text: str, agent_type: str) -> None:
"""Send a prompt and press Enter using the active terminal backend."""
await backend.send_text(session, text)
await asyncio.sleep(_compute_prompt_delay(text, agent_type))
await backend.send_key(session, "enter")
async def _close_single_worker(
backend,
session,
session_id: str,
registry: "SessionRegistry",
force: bool = False,
) -> dict:
"""
Close a single worker session.
Internal helper for close_workers. Handles the actual close logic
for one session.
Args:
session: The ManagedSession object
session_id: ID of the session to close
registry: The session registry
backend: Terminal backend used for terminal operations
force: If True, force close even if session is busy
Returns:
Dict with success status and worktree_cleaned flag
"""
# Check if busy
if session.status == SessionStatus.BUSY and not force:
return {
"success": False,
"error": "Session is busy",
"hint": HINTS["session_busy"],
"worktree_cleaned": False,
}
try:
# Send Ctrl+C to interrupt any running operation
await backend.send_key(session.terminal_session, "ctrl-c")
# TODO(rabsef-bicrym): Programmatically time these actions
await asyncio.sleep(1.0)
if session.agent_type == "codex":
# Codex exits via Ctrl+C (may require a second press).
await backend.send_key(session.terminal_session, "ctrl-c")
# TODO(rabsef-bicrym): Programmatically time these actions
await asyncio.sleep(1.0)
else:
# Claude exits via /exit.
await _send_prompt(backend, session.terminal_session, "/exit", session.agent_type)
# TODO(rabsef-bicrym): Programmatically time these actions
await asyncio.sleep(1.0)
# Clean up worktree if exists (keeps branch alive for cherry-picking)
worktree_cleaned = False
if session.worktree_path and session.main_repo_path:
try:
remove_worktree(
repo_path=session.main_repo_path,
worktree_path=session.worktree_path,
)
worktree_cleaned = True
except WorktreeError as e:
# Log but don't fail the close
logger.warning(f"Failed to clean up worktree for {session_id}: {e}")
# Close the terminal pane/window
await backend.close_session(session.terminal_session, force=force)
# Remove from registry
registry.remove(session_id)
return {
"success": True,
"worktree_cleaned": worktree_cleaned,
}
except Exception as e:
logger.error(f"Failed to close session {session_id}: {e}")
# Still try to remove from registry
registry.remove(session_id)
return {
"success": True,
"warning": f"Session removed but cleanup may be incomplete: {e}",
"worktree_cleaned": False,
}
def register_tools(mcp: FastMCP) -> None:
"""Register close_workers tool on the MCP server."""
@mcp.tool()
async def close_workers(
ctx: Context[ServerSession, "AppContext"],
session_ids: list[str],
force: bool | None = False,
) -> dict:
"""
Close one or more managed Claude Code sessions.
Gracefully terminates the Claude sessions in parallel and closes
their terminal panes/windows. All session_ids must exist in the registry.
⚠️ **NOTE: WORKTREE CLEANUP**
Workers with worktrees commit to ephemeral branches. When closed:
- The worktree directory is removed
- The branch is KEPT for cherry-picking/merging
**AFTER closing workers with worktrees:**
1. Review commits on the worker's branch
2. Merge or cherry-pick commits to a persistent branch
3. Delete the branch when done: `git branch -D <branch-name>`
Args:
session_ids: List of session IDs to close (1 or more required).
Accepts internal IDs, terminal IDs, or worker names.
force: If True, force close even if sessions are busy
Returns:
Dict with:
- session_ids: List of session IDs that were requested
- results: Dict mapping session_id to individual result
- success_count: Number of sessions closed successfully
- failure_count: Number of sessions that failed to close
"""
# Handle None values from MCP clients that send explicit null for omitted params
force = force if force is not None else False
app_ctx = ctx.request_context.lifespan_context
registry = app_ctx.registry
backend = app_ctx.terminal_backend
if not session_ids:
return error_response(
"No session_ids provided",
hint="Provide at least one session_id to close",
)
# Validate all sessions exist first (fail fast)
sessions_to_close = []
missing_sessions = []
for sid in session_ids:
session = registry.resolve(sid)
if not session:
missing_sessions.append(sid)
else:
sessions_to_close.append((sid, session))
# If any sessions are missing, fail the entire operation
if missing_sessions:
return error_response(
f"Sessions not found: {', '.join(missing_sessions)}",
hint=HINTS["session_not_found"],
session_ids=session_ids,
missing=missing_sessions,
)
# Close all sessions in parallel
async def close_one(sid: str, session) -> tuple[str, dict]:
result = await _close_single_worker(backend, session, sid, registry, force)
return (sid, result)
tasks = [close_one(sid, session) for sid, session in sessions_to_close]
parallel_results = await asyncio.gather(*tasks, return_exceptions=True)
# Aggregate results
results = {}
for item in parallel_results:
if isinstance(item, BaseException):
# Shouldn't happen since _close_single_worker catches exceptions
logger.error(f"Unexpected exception in close_workers: {item}")
continue
# Type narrowing: item is now tuple[str, dict]
sid, result = item
results[sid] = result
success_count = sum(1 for r in results.values() if r.get("success", False))
failure_count = len(results) - success_count
return {
"session_ids": session_ids,
"results": results,
"success_count": success_count,
"failure_count": failure_count,
}
