diff --git a/features/steps/tdd_mcp_client_timer_cancel_race_steps.py b/features/steps/tdd_mcp_client_timer_cancel_race_steps.py index c78f45185..11bba652e 100644 --- a/features/steps/tdd_mcp_client_timer_cancel_race_steps.py +++ b/features/steps/tdd_mcp_client_timer_cancel_race_steps.py @@ -42,6 +42,19 @@ def _mock_tool(name: str, desc: str = "", schema: dict | None = None) -> dict: } +class _ShutdownOnCancelTimer: + """Timer probe that marks shutdown while scheduler cancellation is active.""" + + def __init__(self, client: McpClient) -> None: + self._client = client + self.cancel_called = False + + def cancel(self) -> None: + self.cancel_called = True + with self._client._lock: + self._client._shutting_down = True + + # ── Given ───────────────────────────────────────────────────────── @@ -117,6 +130,18 @@ def step_mark_timer_race_client_shutting_down(context: Context) -> None: context.timer_race_client._shutting_down = True +@given("cancellation of existing timers begins shutdown") +def step_existing_timer_cancellation_begins_shutdown(context: Context) -> None: + """Install cancel probes that flip shutdown after the scheduler's first guard.""" + client = context.timer_race_client + context.timer_race_idle_cancel_probe = _ShutdownOnCancelTimer(client) + context.timer_race_health_cancel_probe = _ShutdownOnCancelTimer(client) + with client._lock: + client._idle_timer = context.timer_race_idle_cancel_probe + client._health_timer = context.timer_race_health_cancel_probe + client._shutting_down = False + + # ── When ────────────────────────────────────────────────────────── @@ -245,8 +270,12 @@ def step_wait_for_late_fires(context: Context) -> None: @when("I ask the client to schedule idle and health timers") def step_schedule_idle_and_health_timers(context: Context) -> None: """Invoke both scheduler guards while shutdown is in progress.""" - context.timer_race_client._schedule_idle_timer() - context.timer_race_client._schedule_health_check() + client = context.timer_race_client + client._schedule_idle_timer() + if hasattr(context, "timer_race_health_cancel_probe"): + with client._lock: + client._shutting_down = False + client._schedule_health_check() # ── Then ────────────────────────────────────────────────────────── @@ -281,3 +310,14 @@ def step_no_idle_or_health_timer_scheduled(context: Context) -> None: client = context.timer_race_client assert client._idle_timer is None, "Idle timer was scheduled during shutdown" assert client._health_timer is None, "Health timer was scheduled during shutdown" + + +@then("cancellation should have been observed for both timers") +def step_cancellation_observed_for_both_timers(context: Context) -> None: + """Assert both scheduler paths reached the cancellation probe.""" + assert context.timer_race_idle_cancel_probe.cancel_called, ( + "Idle timer cancellation probe was not reached" + ) + assert context.timer_race_health_cancel_probe.cancel_called, ( + "Health timer cancellation probe was not reached" + ) diff --git a/features/tdd_mcp_client_timer_cancel_race.feature b/features/tdd_mcp_client_timer_cancel_race.feature index 63ac4261e..09e7dc997 100644 --- a/features/tdd_mcp_client_timer_cancel_race.feature +++ b/features/tdd_mcp_client_timer_cancel_race.feature @@ -82,3 +82,10 @@ Feature: TDD Issue #10516 — McpClient timer can fire after cancellation race And the timer race client is marked as shutting down When I ask the client to schedule idle and health timers Then no idle or health timer should be scheduled + + Scenario: Timer schedulers recheck shutdown after cancellation + Given an McpClient with active idle and health timers + And cancellation of existing timers begins shutdown + When I ask the client to schedule idle and health timers + Then cancellation should have been observed for both timers + And no idle or health timer should be scheduled diff --git a/src/cleveragents/mcp/client.py b/src/cleveragents/mcp/client.py index 39c28c602..c83146e08 100644 --- a/src/cleveragents/mcp/client.py +++ b/src/cleveragents/mcp/client.py @@ -330,13 +330,14 @@ class McpClient: with self._lock: if self._shutting_down: return - - self._cancel_idle_timer() - timeout = self._config.idle_timeout_seconds - if timeout <= 0: - return - - with self._lock: + if self._idle_timer is not None: + self._idle_timer.cancel() + self._idle_timer = None + if self._shutting_down: + return + timeout = self._config.idle_timeout_seconds + if timeout <= 0: + return timer = threading.Timer(timeout, self._check_idle) timer.daemon = True self._idle_timer = timer @@ -378,13 +379,14 @@ class McpClient: with self._lock: if self._shutting_down: return - - self._cancel_health_check() - interval = self._config.health_check_interval_seconds - if interval <= 0: - return - - with self._lock: + if self._health_timer is not None: + self._health_timer.cancel() + self._health_timer = None + if self._shutting_down: + return + interval = self._config.health_check_interval_seconds + if interval <= 0: + return timer = threading.Timer(interval, self._perform_health_check) timer.daemon = True self._health_timer = timer