From b55be3f6bdd0bfeeb0220a42bc9017f92a430b0c Mon Sep 17 00:00:00 2001
From: "Slai.life" <huong8373tt@gmail.com>
Date: Wed, 28 Oct 2026 18:18:17 +0800
Subject: [PATCH] Use an absolute deadline for Docker invocations

---
 pyrit/_code_execution/docker_client.py        |  3 +-
 .../_code_execution/test_docker_client.py     | 33 +++++++++++++++++++
 2 files changed, 35 insertions(+), 1 deletion(-)

diff --git a/pyrit/_code_execution/docker_client.py b/pyrit/_code_execution/docker_client.py
index 67ec96763b..435d5c8adf 100644
--- a/pyrit/_code_execution/docker_client.py
+++ b/pyrit/_code_execution/docker_client.py
@@ -386,7 +386,8 @@ async def _run_invocation_async(
         cleanup_warning: str | None = None
         try:
             try:
-                process = await asyncio.wait_for(asyncio.shield(startup), timeout=self._timeout_seconds)
+                async with asyncio.timeout_at(deadline):
+                    process = await asyncio.shield(startup)
             except TimeoutError as error:
                 raise SandboxExecutionError("the Docker client did not finish starting before its deadline") from error
             result = await self._wait_for_invocation_async(
diff --git a/tests/unit/_code_execution/test_docker_client.py b/tests/unit/_code_execution/test_docker_client.py
index 9f65d8219b..5effabe9a9 100644
--- a/tests/unit/_code_execution/test_docker_client.py
+++ b/tests/unit/_code_execution/test_docker_client.py
@@ -16,6 +16,7 @@
 import pytest
 
 from pyrit._code_execution import DockerRuntimeProfile, DockerSandboxClient, SandboxExecutionError
+from pyrit._code_execution.docker_client import _ClientOutput, _ContainerState, _InvocationResult
 from pyrit.models.code_execution import _ProgramInvocation, _SandboxRun
 
 if TYPE_CHECKING:
@@ -236,6 +237,38 @@ async def test_container_and_client_outcomes_async(
     assert await asyncio.to_thread(lambda: list(tmp_path.iterdir())) == []
 
 
+async def test_startup_consumes_shared_deadline_before_capture_async(client: DockerSandboxClient) -> None:
+    process = _process()
+    scope_entered = asyncio.Event()
+
+    async def start_async(*, argv: list[str]) -> MagicMock:
+        del argv
+        assert scope_entered.is_set()
+        return process
+
+    output = _ClientOutput(stdout=b"", stdout_truncated=False, stderr=b"", stderr_truncated=False)
+    container = _ContainerState(status="exited", exit_code=0, oom_killed=False)
+    wait = AsyncMock(return_value=_InvocationResult(output=output, container=container))
+
+    class _NoopTimeout:
+        async def __aenter__(self) -> None:
+            scope_entered.set()
+
+        async def __aexit__(self, *args: object) -> bool:
+            return False
+
+    with (
+        patch("pyrit._code_execution.docker_client.asyncio.timeout_at", return_value=_NoopTimeout()) as timeout_at,
+        patch.object(client, "_start_client_async", side_effect=start_async),
+        patch.object(client, "_wait_for_invocation_async", wait),
+        patch.object(client, "_control_async", _control()),
+    ):
+        assert (await _run_async(client)).events[0].exit_code == 0
+
+    deadline = timeout_at.call_args.args[0]
+    assert wait.call_args.kwargs["deadline"] == deadline
+
+
 @pytest.mark.parametrize(
     ("eof", "state", "error"),
     [