fix: address copilot feedback

Signed-off-by: Samantha Coyle <sam@diagrid.io>

Author: sicoyle

ext/dapr-ext-workflow/dapr/ext/workflow/aio/mcp.py

@@ -102,17 +102,29 @@ async def connect(self, mcpserver_name: str) -> None:
                 )
                 break
             except Exception as exc:  # noqa: BLE001 — classified by helper
-                if not _is_transient_schedule_error(exc) or time.monotonic() >= deadline:
+                if not _is_transient_schedule_error(exc):
+                    raise
+                sleep_for = min(
+                    _SCHEDULE_RETRY_INTERVAL_SECONDS, deadline - time.monotonic()
+                )
+                if sleep_for <= 0:
                     raise
                 logger.debug(
                     'schedule_new_workflow returned transient error %s; retrying', exc
                 )
-                await asyncio.sleep(_SCHEDULE_RETRY_INTERVAL_SECONDS)
+                await asyncio.sleep(sleep_for)
 
-        remaining = max(deadline - time.monotonic(), 1.0)
+        remaining = deadline - time.monotonic()
+        if remaining <= 0:
+            raise RuntimeError(
+                f"ListTools workflow for MCPServer '{mcpserver_name}' "
+                f'timed out after {self._timeout}s'
+            )
+        # wait_for_workflow_completion treats timeout=0 as "wait forever",
+        # so floor the gRPC timeout at 1s when sub-second remaining survives.
         state = await self._wf_client.wait_for_workflow_completion(
             instance_id=instance_id,
-            timeout_in_seconds=int(remaining),
+            timeout_in_seconds=max(int(remaining), 1),
             fetch_payloads=True,
         )
 

ext/dapr-ext-workflow/dapr/ext/workflow/mcp.py

@@ -245,17 +245,29 @@ def connect(self, mcpserver_name: str) -> None:
                 )
                 break
             except Exception as exc:  # noqa: BLE001 — classified by helper
-                if not _is_transient_schedule_error(exc) or time.monotonic() >= deadline:
+                if not _is_transient_schedule_error(exc):
+                    raise
+                sleep_for = min(
+                    _SCHEDULE_RETRY_INTERVAL_SECONDS, deadline - time.monotonic()
+                )
+                if sleep_for <= 0:
                     raise
                 logger.debug(
                     'schedule_new_workflow returned transient error %s; retrying', exc
                 )
-                time.sleep(_SCHEDULE_RETRY_INTERVAL_SECONDS)
+                time.sleep(sleep_for)
 
-        remaining = max(deadline - time.monotonic(), 1.0)
+        remaining = deadline - time.monotonic()
+        if remaining <= 0:
+            raise RuntimeError(
+                f"ListTools workflow for MCPServer '{mcpserver_name}' "
+                f'timed out after {self._timeout}s'
+            )
+        # wait_for_workflow_completion treats timeout=0 as "wait forever",
+        # so floor the gRPC timeout at 1s when sub-second remaining survives.
         state = self._wf_client.wait_for_workflow_completion(
             instance_id=instance_id,
-            timeout_in_seconds=int(remaining),
+            timeout_in_seconds=max(int(remaining), 1),
             fetch_payloads=True,
         )
 

ext/dapr-ext-workflow/tests/test_mcp_client.py

@@ -469,6 +469,28 @@ def test_deadline_exhausted_raises_last_error(self):
             with self.assertRaises(grpc.RpcError):
                 mcp_client.connect('weather')
 
+    def test_budget_exhausted_after_schedule_succeeds(self):
+        """If retries burn the budget but schedule eventually succeeds, raise
+        without calling wait_for_workflow_completion (timeout=0 means
+        'wait forever' in the underlying client)."""
+        mock_wf = MagicMock()
+        mock_wf.schedule_new_workflow.side_effect = [
+            _StubRpcError(grpc.StatusCode.CANCELLED),
+            'inst-1',
+        ]
+
+        mcp_client = DaprMCPClient(timeout_in_seconds=1, wf_client=mock_wf)
+        # monotonic: 0.0 → deadline = 1.0; 0.4 → sleep_for = 0.5 (still in budget);
+        # 2.0 → post-loop remaining = -1.0 → raise.
+        with patch('dapr.ext.workflow.mcp.time.sleep'), patch(
+            'dapr.ext.workflow.mcp.time.monotonic',
+            side_effect=[0.0, 0.4, 2.0],
+        ):
+            with self.assertRaises(RuntimeError) as ctx:
+                mcp_client.connect('weather')
+        self.assertIn('timed out', str(ctx.exception))
+        mock_wf.wait_for_workflow_completion.assert_not_called()
+
 
 class TestAioDaprMCPClientConnectRetry(unittest.IsolatedAsyncioTestCase):
     """Async counterpart of TestDaprMCPClientConnectRetry."""
@@ -504,6 +526,24 @@ async def test_deadline_exhausted_raises(self):
             with self.assertRaises(grpc.RpcError):
                 await mcp_client.connect('weather')
 
+    async def test_budget_exhausted_after_schedule_succeeds(self):
+        """Async mirror of the fail-fast-after-schedule-success guard."""
+        mock_wf = AsyncMock()
+        mock_wf.schedule_new_workflow.side_effect = [
+            _StubRpcError(grpc.StatusCode.CANCELLED),
+            'inst-1',
+        ]
+
+        mcp_client = AioDaprMCPClient(timeout_in_seconds=1, wf_client=mock_wf)
+        with patch('dapr.ext.workflow.aio.mcp.asyncio.sleep', new=AsyncMock()), patch(
+            'dapr.ext.workflow.aio.mcp.time.monotonic',
+            side_effect=[0.0, 0.4, 2.0],
+        ):
+            with self.assertRaises(RuntimeError) as ctx:
+                await mcp_client.connect('weather')
+        self.assertIn('timed out', str(ctx.exception))
+        mock_wf.wait_for_workflow_completion.assert_not_awaited()
+
 
 class TestMCPWorkflowPrefix(unittest.TestCase):
     """Tests for the workflow naming constant."""