From c0a96a39c6feb06270a6b2d45404f782a15f87b7 Mon Sep 17 00:00:00 2001
From: TJ Bruno <tj.bruno@tensordyne.ai>
Date: Fri, 10 Apr 2026 16:30:53 -0700
Subject: [PATCH 1/3] fix: Apply gevent timeout to function command operation
 call

---
 src/pyinfra/api/operations.py | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/src/pyinfra/api/operations.py b/src/pyinfra/api/operations.py
index 5bca320d3..17bdf0e95 100644
--- a/src/pyinfra/api/operations.py
+++ b/src/pyinfra/api/operations.py
@@ -103,7 +103,12 @@ def _run_host_op(state: "State", host: "Host", op_hash: str) -> bool:
 
             if isinstance(command, FunctionCommand):
                 try:
-                    status = command.execute(state, host, connector_arguments)
+                    with gevent.Timeout(timeout, exception=TimeoutError):
+                        status = command.execute(state, host, connector_arguments)
+
+                except TimeoutError as e:
+                    log_host_command_error(host, e, timeout=timeout)
+
                 except NestedOperationError:
                     host.log_styled("Error in nested operation", fg="red", log_func=logger.error)
                 except Exception as e:

From c0c4c3c7d856a8d1c9ac6ba1ce60a6c3d4efc86c Mon Sep 17 00:00:00 2001
From: TJ Bruno <tj.bruno@tensordyne.ai>
Date: Fri, 10 Apr 2026 16:31:48 -0700
Subject: [PATCH 2/3] fix: Log function command timeouts

---
 src/pyinfra/api/util.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/pyinfra/api/util.py b/src/pyinfra/api/util.py
index aaf565eef..e10a20b1f 100644
--- a/src/pyinfra/api/util.py
+++ b/src/pyinfra/api/util.py
@@ -261,7 +261,7 @@ def log_error_or_warning(
 
 
 def log_host_command_error(host: "Host", e: Exception, timeout: int | None = 0) -> None:
-    if isinstance(e, timeout_error):
+    if isinstance(e, (TimeoutError, timeout_error)):
         logger.error(
             "{0}{1}".format(
                 host.print_prefix,

From eabad46e9384797fb19bfb19ea710cf82eb54aad Mon Sep 17 00:00:00 2001
From: TJ Bruno <tj.bruno@tensordyne.ai>
Date: Fri, 10 Apr 2026 17:09:20 -0700
Subject: [PATCH 3/3] test: Add unit test coverage

---
 tests/test_api/test_api_operations.py | 18 ++++++++++++++++++
 1 file changed, 18 insertions(+)

diff --git a/tests/test_api/test_api_operations.py b/tests/test_api/test_api_operations.py
index b3edd3a1a..d1cbb645c 100644
--- a/tests/test_api/test_api_operations.py
+++ b/tests/test_api/test_api_operations.py
@@ -2,6 +2,7 @@
 from os import path
 from unittest import TestCase
 from unittest.mock import mock_open, patch
+import time
 
 import pyinfra
 from pyinfra.api import (
@@ -258,6 +259,23 @@ def mocked_function(*args, **kwargs):
 
         assert is_called
 
+    def test_function_call_op_timeout(self):
+        inventory = make_inventory()
+        state = State(inventory, Config())
+        state.current_stage = StateStage.Prepare
+        connect_all(state)
+
+        timeout = 1
+
+        def mocked_function(*args, **kwargs):
+            time.sleep(timeout + 1)
+
+        add_op(state, python.call, mocked_function, _timeout=timeout)
+
+        # Timeout should cause the operation to fail and hosts to be removed
+        with self.assertRaises(PyinfraError) as context:
+            run_ops(state)
+
     def test_run_once_serial_op(self):
         inventory = make_inventory()
         state = State(inventory, Config())