diff --git a/backend/onyx/server/features/build/sandbox/base.py b/backend/onyx/server/features/build/sandbox/base.py index 6926c5cf88..b2aa669ad2 100644 --- a/backend/onyx/server/features/build/sandbox/base.py +++ b/backend/onyx/server/features/build/sandbox/base.py @@ -349,6 +349,7 @@ class SandboxManager(_ServeMixin, ABC): on_opencode_session_resolved: Callable[[str], None] | None = None, should_interrupt: Callable[[], bool] | None = None, should_abort_on_teardown: Callable[[], bool] | None = None, + turn_timeout_seconds: float | None = None, ) -> Generator[SandboxEvent, None, None]: """Stream typed sandbox events for one user message via opencode-serve. @@ -371,6 +372,7 @@ class SandboxManager(_ServeMixin, ABC): on_opencode_session_resolved=on_opencode_session_resolved, should_interrupt=should_interrupt, should_abort_on_teardown=should_abort_on_teardown, + turn_timeout_seconds=turn_timeout_seconds, ) def send_subagent_message( diff --git a/backend/onyx/server/features/build/sandbox/serve_transport.py b/backend/onyx/server/features/build/sandbox/serve_transport.py index 012ec1813c..2280d487d6 100644 --- a/backend/onyx/server/features/build/sandbox/serve_transport.py +++ b/backend/onyx/server/features/build/sandbox/serve_transport.py @@ -27,6 +27,7 @@ from onyx.cache.interface import CACHE_TRANSIENT_ERRORS from onyx.cache.interface import CacheLock from onyx.cache.interface import CacheLockLostError from onyx.db.engine.sql_engine import get_session_with_current_tenant +from onyx.server.features.build.configs import OPENCODE_PROMPT_TIMEOUT_SECONDS from onyx.server.features.build.configs import OPENCODE_SERVE_EVENT_READ_TIMEOUT from onyx.server.features.build.configs import OPENCODE_SERVER_USERNAME from onyx.server.features.build.configs import PROMPT_SLOT_LEASE_SECONDS @@ -541,6 +542,7 @@ class _ServeMixin: on_opencode_session_resolved: Callable[[str], None] | None = None, should_interrupt: Callable[[], bool] | None = None, should_abort_on_teardown: Callable[[], bool] | None = None, + turn_timeout_seconds: float | None = None, ) -> Generator[SandboxEvent, None, None]: """Stream sandbox events via the in-sandbox ``opencode serve``. Preflight ``opencode_session_id`` via :meth:`ensure_opencode_session` to avoid @@ -589,6 +591,11 @@ class _ServeMixin: directory=session_path, model_provider=agent_provider, model_id=agent_model, + timeout=( + turn_timeout_seconds + if turn_timeout_seconds is not None + else OPENCODE_PROMPT_TIMEOUT_SECONDS + ), should_interrupt=should_interrupt, ): events_count += 1 diff --git a/backend/onyx/server/features/build/scheduled_tasks/executor.py b/backend/onyx/server/features/build/scheduled_tasks/executor.py index e498a02123..5b617563c9 100644 --- a/backend/onyx/server/features/build/scheduled_tasks/executor.py +++ b/backend/onyx/server/features/build/scheduled_tasks/executor.py @@ -419,7 +419,10 @@ def _drive_agent( return False try: for sandbox_event in session_manager.yield_sandbox_events( - sandbox_id, session_id, task_prompt + sandbox_id, + session_id, + task_prompt, + turn_timeout_seconds=float(budget_seconds), ): slot.extend() if slot.lost: diff --git a/backend/onyx/server/features/build/session/manager.py b/backend/onyx/server/features/build/session/manager.py index eb8b911d43..1000cfb139 100644 --- a/backend/onyx/server/features/build/session/manager.py +++ b/backend/onyx/server/features/build/session/manager.py @@ -838,6 +838,7 @@ class SessionManager: user_message_content: str, should_interrupt: Callable[[], bool] | None = None, should_abort_on_teardown: Callable[[], bool] | None = None, + turn_timeout_seconds: float | None = None, ) -> Generator[Any, None, None]: build_session = _streaming.load_turn_session( self._db_session, self._sandbox_manager, sandbox_id, session_id @@ -855,6 +856,7 @@ class SessionManager: agent_model=build_session.agent_model, should_interrupt=should_interrupt, should_abort_on_teardown=should_abort_on_teardown, + turn_timeout_seconds=turn_timeout_seconds, ) def merge_events_with_announces( diff --git a/backend/onyx/server/features/build/session/streaming.py b/backend/onyx/server/features/build/session/streaming.py index 9a57265046..5eb5764870 100644 --- a/backend/onyx/server/features/build/session/streaming.py +++ b/backend/onyx/server/features/build/session/streaming.py @@ -486,6 +486,7 @@ def yield_sandbox_events( agent_model: str | None, should_interrupt: Callable[[], bool] | None = None, should_abort_on_teardown: Callable[[], bool] | None = None, + turn_timeout_seconds: float | None = None, ) -> Generator[Any, None, None]: """Drive the agent to completion, yielding raw sandbox events. @@ -520,6 +521,7 @@ def yield_sandbox_events( on_opencode_session_resolved=_persist_resolved_id, should_interrupt=should_interrupt, should_abort_on_teardown=should_abort_on_teardown, + turn_timeout_seconds=turn_timeout_seconds, ) try: for sandbox_event in event_stream: diff --git a/backend/tests/common/craft/stubs.py b/backend/tests/common/craft/stubs.py index 05f60a20c5..a97b34ad9f 100644 --- a/backend/tests/common/craft/stubs.py +++ b/backend/tests/common/craft/stubs.py @@ -475,6 +475,7 @@ class StubSandboxManager(SandboxManager): on_opencode_session_resolved: Callable[[str], None] | None = None, should_interrupt: Callable[[], bool] | None = None, should_abort_on_teardown: Callable[[], bool] | None = None, + turn_timeout_seconds: float | None = None, ) -> Generator[SandboxEvent, None, None]: self.send_message_count += 1 self.last_send_message_payload = { @@ -487,6 +488,7 @@ class StubSandboxManager(SandboxManager): "on_opencode_session_resolved": on_opencode_session_resolved, "should_interrupt": should_interrupt, "should_abort_on_teardown": should_abort_on_teardown, + "turn_timeout_seconds": turn_timeout_seconds, } if self._send_message_events is None: raise _not_configured("send_message") diff --git a/backend/tests/external_dependency_unit/craft/test_scheduled_task_executor.py b/backend/tests/external_dependency_unit/craft/test_scheduled_task_executor.py index 3c0d6cdd45..784cdebfff 100644 --- a/backend/tests/external_dependency_unit/craft/test_scheduled_task_executor.py +++ b/backend/tests/external_dependency_unit/craft/test_scheduled_task_executor.py @@ -367,6 +367,44 @@ def test_prompt_response_marks_run_succeeded( assert stub_sandbox_manager.last_prompt_slot.extend_calls >= 1 +def test_scheduled_run_threads_budget_as_turn_timeout( + db_session: Session, + test_user: User, + sandbox: Any, # noqa: ARG001 + session_manager_with_stub: SessionManager, # noqa: ARG001 + stub_sandbox_manager: StubSandboxManager, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """budget_seconds must thread through as the serve client's per-turn + timeout so the generic 15-min prompt timeout can't undercut the run budget.""" + # Bypass skill-payload: encrypted ExternalApp creds break local MIT decryption. + monkeypatch.setattr( + "onyx.server.features.build.session.manager.build_user_skills_payload", + lambda *_: ("", "", {}), + ) + + sandbox(user=test_user, status=SandboxStatus.RUNNING) + _, run = _seed_task_and_queued_run(db_session, test_user) + + stub_sandbox_manager.health_check_returns = True + stub_sandbox_manager.setup_session_workspace_silent = True + stub_sandbox_manager.write_files_to_sandbox_silent = True + stub_sandbox_manager.send_message_events = [ + PromptResponse.model_validate({"stopReason": "end_turn"}), + ] + + run_scheduled_task_logic(run.id, budget_seconds=1234) + + send_message_payload = stub_sandbox_manager.last_send_message_payload + assert send_message_payload is not None + assert send_message_payload["turn_timeout_seconds"] == 1234.0 + + db_session.expire_all() + refreshed = db_session.get(ScheduledTaskRun, run.id) + assert refreshed is not None + assert refreshed.status == ScheduledTaskRunStatus.SUCCEEDED + + def test_cancelled_prompt_response_marks_run_failed( db_session: Session, test_user: User, diff --git a/backend/tests/unit/onyx/server/features/build/sandbox/test_send_message_timeout_plumbing.py b/backend/tests/unit/onyx/server/features/build/sandbox/test_send_message_timeout_plumbing.py new file mode 100644 index 0000000000..14f302e719 --- /dev/null +++ b/backend/tests/unit/onyx/server/features/build/sandbox/test_send_message_timeout_plumbing.py @@ -0,0 +1,78 @@ +"""`turn_timeout_seconds` must reach `OpencodeServeClient.send_message` as +`timeout`, defaulting to `OPENCODE_PROMPT_TIMEOUT_SECONDS` when unset.""" + +from __future__ import annotations + +from collections.abc import Callable +from collections.abc import Generator +from typing import Any +from uuid import uuid4 + +from onyx.server.features.build.configs import OPENCODE_PROMPT_TIMEOUT_SECONDS +from onyx.server.features.build.sandbox.event_schema import PromptResponse +from onyx.server.features.build.sandbox.kubernetes.kubernetes_sandbox_manager import ( + KubernetesSandboxManager, +) + + +class _FakeServeClient: + def __init__(self) -> None: + self.captured_timeout: float | None = None + + def ensure_session( + self, + opencode_session_id: str | None, # noqa: ARG002 + *, + directory: str, # noqa: ARG002 + title: str | None = None, # noqa: ARG002 + ) -> str: + return "ses_fake" + + def send_message( + self, + opencode_session_id: str, # noqa: ARG002 + message: str, # noqa: ARG002 + *, + directory: str, # noqa: ARG002 + model_provider: str | None = None, # noqa: ARG002 + model_id: str | None = None, # noqa: ARG002 + timeout: float, + should_interrupt: Callable[[], bool] | None = None, # noqa: ARG002 + ) -> Generator[Any, None, None]: + self.captured_timeout = timeout + yield PromptResponse.model_validate({"stopReason": "end_turn"}) + + def close(self) -> None: + pass + + +def _manager_with(client: _FakeServeClient) -> KubernetesSandboxManager: + manager: KubernetesSandboxManager = object.__new__(KubernetesSandboxManager) + manager._init_serve_state() + + def build_client(*_: Any, **__: Any) -> _FakeServeClient: + return client + + manager._build_serve_client = build_client # type: ignore[method-assign] + return manager + + +def test_send_message_threads_turn_timeout_to_client() -> None: + client = _FakeServeClient() + manager = _manager_with(client) + + events = list( + manager.send_message(uuid4(), uuid4(), "hi", turn_timeout_seconds=1234.0) + ) + + assert client.captured_timeout == 1234.0 + assert any(isinstance(e, PromptResponse) for e in events) + + +def test_send_message_defaults_to_prompt_timeout() -> None: + client = _FakeServeClient() + manager = _manager_with(client) + + list(manager.send_message(uuid4(), uuid4(), "hi")) + + assert client.captured_timeout == OPENCODE_PROMPT_TIMEOUT_SECONDS diff --git a/backend/tests/unit/onyx/server/features/build/session/test_streaming_state.py b/backend/tests/unit/onyx/server/features/build/session/test_streaming_state.py index aa17b7a7d7..5eb5e663d8 100644 --- a/backend/tests/unit/onyx/server/features/build/session/test_streaming_state.py +++ b/backend/tests/unit/onyx/server/features/build/session/test_streaming_state.py @@ -39,6 +39,7 @@ class _FakeStreamingSandboxManager: on_opencode_session_resolved: Any = None, # noqa: ARG002 should_interrupt: Any = None, # noqa: ARG002 should_abort_on_teardown: Any = None, # noqa: ARG002 + turn_timeout_seconds: float | None = None, # noqa: ARG002 ) -> Any: self.last_payload = { "sandbox_id": sandbox_id,