From 1997100a569e980488d92bda438caadb87052720 Mon Sep 17 00:00:00 2001
From: Wylabb <77673282+Wylabb@users.noreply.github.com>
Date: Sun, 30 Aug 2026 11:02:02 +0200
Subject: [PATCH] fix: treat prompt timeout as inactivity limit

---
 src/prime_agent_client/session.py | 11 ++++++-----
 tests/fixtures/fake_prime_rpc.py  |  8 ++++++++
 tests/test_client.py              | 13 +++++++++++++
 3 files changed, 27 insertions(+), 5 deletions(-)

diff --git a/src/prime_agent_client/session.py b/src/prime_agent_client/session.py
index 6c7cede..950613c 100644
--- a/src/prime_agent_client/session.py
+++ b/src/prime_agent_client/session.py
@@ -167,15 +167,16 @@ async def prompt_stream(
     ) -> AsyncIterator[PrimeEvent]:
         stream = self.events()
         timeout_seconds = self.prompt_timeout if timeout is None else float(timeout)
-        deadline = asyncio.get_running_loop().time() + timeout_seconds
         try:
             await self.prompt(message, images=images)
             while True:
                 try:
-                    remaining = deadline - asyncio.get_running_loop().time()
-                    if remaining <= 0:
-                        raise asyncio.TimeoutError
-                    event = await asyncio.wait_for(stream.__anext__(), timeout=remaining)
+                    # Agent turns can legitimately run for hours. Treat the prompt
+                    # timeout as an inactivity limit, not a wall-clock deadline, so
+                    # a live event stream is never aborted merely for being long.
+                    event = await asyncio.wait_for(
+                        stream.__anext__(), timeout=timeout_seconds
+                    )
                 except asyncio.TimeoutError as exc:
                     with contextlib.suppress(Exception):
                         await self.abort()
diff --git a/tests/fixtures/fake_prime_rpc.py b/tests/fixtures/fake_prime_rpc.py
index ecb91da..427441f 100644
--- a/tests/fixtures/fake_prime_rpc.py
+++ b/tests/fixtures/fake_prime_rpc.py
@@ -5,6 +5,7 @@
 import json
 import os
 import sys
+import time
 from typing import Any
 
 
@@ -87,6 +88,13 @@ def emit(payload: dict[str, Any]) -> None:
     if command == "prompt":
         if request.get("message") == "stall":
             continue
+        if request.get("message") == "active-beyond-timeout":
+            emit({"type": "agent_start"})
+            for index in range(3):
+                time.sleep(0.03)
+                emit({"type": "progress", "index": index})
+            emit({"type": "agent_end"})
+            continue
         if request.get("message") == "ui":
             waiting_for_ui = True
             emit({"type": "extension_ui_request", "id": "ui-1", "method": "confirm"})
diff --git a/tests/test_client.py b/tests/test_client.py
index 1f8eac8..f1a91d6 100644
--- a/tests/test_client.py
+++ b/tests/test_client.py
@@ -215,3 +215,16 @@ async def test_prompt_stream_timeout_aborts_the_active_run() -> None:
             await session.prompt_and_wait("stall")
 
     assert raised.value.command == "prompt_events"
+
+
+async def test_prompt_stream_timeout_measures_inactivity_not_total_runtime() -> None:
+    async with PrimeSession(command=FAKE_COMMAND, prompt_timeout=0.05) as session:
+        events = await session.prompt_and_wait("active-beyond-timeout")
+
+    assert [event.type for event in events] == [
+        "agent_start",
+        "progress",
+        "progress",
+        "progress",
+        "agent_end",
+    ]