diff --git a/plugins/Runner/LocalAgent/components/runner/default.py b/plugins/Runner/LocalAgent/components/runner/default.py index 58ecbd39a..8edb2a54e 100644 --- a/plugins/Runner/LocalAgent/components/runner/default.py +++ b/plugins/Runner/LocalAgent/components/runner/default.py @@ -49,12 +49,18 @@ class RunCancelledError(Exception): class RunDeadline: """Wall-clock deadline shared by assembly, model calls, and tool calls.""" - def __init__(self, timeout_seconds: float): + def __init__(self, timeout_seconds: float, runtime: Any = None): self.timeout_seconds = max(0.001, float(timeout_seconds)) self._deadline = time.monotonic() + self.timeout_seconds + self._runtime = runtime + self._initial_host_deadline = getattr(runtime, "deadline_at", None) def remaining(self) -> float: - return max(0.0, self._deadline - time.monotonic()) + extension = 0.0 + current_host_deadline = getattr(self._runtime, "deadline_at", None) + if self._initial_host_deadline is not None and current_host_deadline is not None: + extension = float(current_host_deadline) - float(self._initial_host_deadline) + return max(0.0, self._deadline + extension - time.monotonic()) class RunInterruptChecker: @@ -148,16 +154,19 @@ class RunInterruptChecker: cancellation = asyncio.create_task(self._cancelled.wait()) temporary_watcher = asyncio.create_task(self._watch()) if self.available and self._watcher is None else None try: - await asyncio.wait( - {task, cancellation}, - timeout=deadline.remaining() if deadline else None, - return_when=asyncio.FIRST_COMPLETED, - ) - if self._cancelled.is_set(): - raise RunCancelledError(CANCELLED_ERROR) - if not task.done(): - raise asyncio.TimeoutError(TIMEOUT_ERROR) - return task.result() + while True: + if self._cancelled.is_set(): + raise RunCancelledError(CANCELLED_ERROR) + if task.done(): + return task.result() + remaining = deadline.remaining() if deadline else None + if remaining is not None and remaining <= 0: + raise asyncio.TimeoutError(TIMEOUT_ERROR) + await asyncio.wait( + {task, cancellation}, + timeout=min(remaining, self.interval_seconds) if remaining is not None else None, + return_when=asyncio.FIRST_COMPLETED, + ) finally: pending = [task, cancellation] + ([temporary_watcher] if temporary_watcher else []) for future in pending: @@ -222,7 +231,7 @@ class DefaultRunner(Runner): api = self.get_run_api(ctx) config = ctx.config if isinstance(ctx.config, dict) else {} timeout_seconds = get_run_timeout_seconds(config) - deadline = RunDeadline(timeout_seconds) if timeout_seconds is not None else None + deadline = RunDeadline(timeout_seconds, ctx.runtime) if timeout_seconds is not None else None usage_tracker = RunUsageTracker() if await interrupt_checker.is_cancelled(force=True): diff --git a/plugins/Runner/LocalAgent/tests/test_run_controls.py b/plugins/Runner/LocalAgent/tests/test_run_controls.py index 7b4f28893..ed2b8a6cf 100644 --- a/plugins/Runner/LocalAgent/tests/test_run_controls.py +++ b/plugins/Runner/LocalAgent/tests/test_run_controls.py @@ -3,6 +3,7 @@ from __future__ import annotations import asyncio +import time import pytest @@ -31,6 +32,25 @@ async def test_slow_ledger_does_not_delay_chunks_or_deadline(): assert checker._watcher is None +@pytest.mark.asyncio +async def test_interaction_extension_keeps_local_runner_wait_alive(): + context = make_context() + context.runtime.deadline_at = time.time() + 0.03 + deadline = RunDeadline(0.03, context.runtime) + checker = RunInterruptChecker(None, context, interval_seconds=0.01) + + async def interaction(): + await asyncio.sleep(0.01) + context.runtime.deadline_at += 0.1 + await asyncio.sleep(0.06) + return "submitted" + + assert await checker.wait_for(interaction(), deadline=deadline) == "submitted" + assert deadline.remaining() > 0 + assert await checker.wait_for(interaction(), deadline=deadline) == "submitted" + assert deadline.remaining() > 0 + + @pytest.mark.asyncio async def test_operation_timeout_is_not_mistaken_for_poll_timeout(): checker = RunInterruptChecker(None, make_context()) diff --git a/pyproject.toml b/pyproject.toml index 8264f86f2..760f8d9fa 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -126,7 +126,7 @@ langbot = "langbot.__main__:main" # Replace this source with the released SDK version before merging #2787. [tool.uv.sources] -langbot-plugin = { git = "https://github.com/fdc310/langbot-plugin-sdk.git", rev = "33dfff35859c30c63acc54deacaaf5b71cf6ef5c" } +langbot-plugin = { git = "https://github.com/fdc310/langbot-plugin-sdk.git", rev = "b8be722481698c14b80b1674957eadab53797406" } [build-system] requires = ["setuptools>=61.0", "wheel"] diff --git a/src/langbot/pkg/plugin/handler.py b/src/langbot/pkg/plugin/handler.py index 6206d2ee1..bc19381f7 100644 --- a/src/langbot/pkg/plugin/handler.py +++ b/src/langbot/pkg/plugin/handler.py @@ -3067,7 +3067,6 @@ class RuntimeConnectionHandler(handler.Handler): Yields RunnerResult dicts. """ - timeout = self._get_runner_action_timeout(context) gen = self.call_action_generator( LangBotToRuntimeAction.RUN_RUNNER, { @@ -3076,28 +3075,13 @@ class RuntimeConnectionHandler(handler.Handler): 'runner_name': runner_name, 'context': context, }, - timeout=timeout, + timeout=None, ) async with contextlib.aclosing(gen): async for ret in gen: yield ret - def _get_runner_action_timeout(self, context: dict[str, Any]) -> float: - """Use the run deadline as the transport idle timeout when available.""" - try: - import time - - deadline_at = (context.get('runtime') or {}).get('deadline_at') - if deadline_at is None: - return 300 - remaining = float(deadline_at) - time.time() - if remaining <= 0: - return 0.001 - return max(remaining + 1.0, 0.001) - except (TypeError, ValueError): - return 300 - async def get_plugin_icon(self, plugin_author: str, plugin_name: str) -> dict[str, Any]: """Get plugin icon""" result = await self.call_action( diff --git a/tests/unit_tests/plugin/test_runner_stream_context.py b/tests/unit_tests/plugin/test_runner_stream_context.py index 76df1bad8..dc25d3d55 100644 --- a/tests/unit_tests/plugin/test_runner_stream_context.py +++ b/tests/unit_tests/plugin/test_runner_stream_context.py @@ -68,6 +68,22 @@ async def test_scope_is_reset_before_yield_and_stream_finishes(deadline): assert closed == [TEST_INSTALLATION_BINDING] +@pytest.mark.asyncio +async def test_runner_transport_wait_has_no_stale_idle_deadline(): + stream, handler, _, _, _, _ = make_stream() + original = handler.call_action_generator + timeouts = [] + + async def capture_timeout(*args, **kwargs): + timeouts.append(kwargs['timeout']) + async for frame in original(*args, **kwargs): + yield frame + + handler.call_action_generator = capture_timeout + assert [frame['sequence'] async for frame in stream] == [0, 1, 2] + assert timeouts == [None] + + @pytest.mark.asyncio @pytest.mark.parametrize('deadline', [True, False]) async def test_early_close_releases_wire_stream_in_scope(deadline): diff --git a/uv.lock b/uv.lock index 14c6e78e8..4543d8091 100644 --- a/uv.lock +++ b/uv.lock @@ -2241,7 +2241,7 @@ requires-dist = [ { name = "gewechat-client", specifier = ">=0.1.5" }, { name = "html2text", specifier = ">=2024.2.26" }, { name = "httpx", extras = ["socks"], specifier = ">=0.28.1" }, - { name = "langbot-plugin", git = "https://github.com/fdc310/langbot-plugin-sdk.git?rev=33dfff35859c30c63acc54deacaaf5b71cf6ef5c" }, + { name = "langbot-plugin", git = "https://github.com/fdc310/langbot-plugin-sdk.git?rev=b8be722481698c14b80b1674957eadab53797406" }, { name = "langchain", specifier = ">=1.3.9" }, { name = "langchain-core", specifier = ">=1.3.3" }, { name = "langchain-text-splitters", specifier = ">=1.1.2" }, @@ -2312,7 +2312,7 @@ dev = [ [[package]] name = "langbot-plugin" version = "0.7.19" -source = { git = "https://github.com/fdc310/langbot-plugin-sdk.git?rev=33dfff35859c30c63acc54deacaaf5b71cf6ef5c#33dfff35859c30c63acc54deacaaf5b71cf6ef5c" } +source = { git = "https://github.com/fdc310/langbot-plugin-sdk.git?rev=b8be722481698c14b80b1674957eadab53797406#b8be722481698c14b80b1674957eadab53797406" } dependencies = [ { name = "aiofiles" }, { name = "aiohttp" },