diff --git a/openhands/automation/dispatcher.py b/openhands/automation/dispatcher.py index da648a04..6464076d 100644 --- a/openhands/automation/dispatcher.py +++ b/openhands/automation/dispatcher.py @@ -497,7 +497,18 @@ async def _fail( if result.success: await update_run_current_phase(session_factory, run.id, "Starting automation") if ctx.sandbox_id: - await update_sandbox_id(session_factory, run.id, ctx.sandbox_id) + recorded = await update_sandbox_id(session_factory, run.id, ctx.sandbox_id) + if recorded is False: + # The run left RUNNING while provisioning (cancelled or + # failed concurrently): drop the sandbox instead of + # attaching it to a terminal row nobody will clean up. + logger.info( + "Run %s left RUNNING during provisioning; releasing sandbox", + run_id, + extra=_log_ctx(sandbox_id=ctx.sandbox_id), + ) + await backend.release_context(client, ctx) + return if result.bash_command_id: # Persist the BashCommand id so the verifier can filter # BashOutput events by exactly this command (avoids diff --git a/openhands/automation/router.py b/openhands/automation/router.py index 2fcf3502..98914669 100644 --- a/openhands/automation/router.py +++ b/openhands/automation/router.py @@ -18,7 +18,7 @@ status, ) from fastapi.responses import RedirectResponse -from sqlalchemy import func, select, update +from sqlalchemy import case, func, select, update from sqlalchemy.engine import CursorResult from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.orm import selectinload @@ -1003,6 +1003,16 @@ async def cancel_run( completed_at=now, error_detail="Cancelled by user", status_detail=None, + # A cancelled run is terminal, so it can never open or continue + # its conversation afterwards: release its subject atomically in + # this same update, whether or not it ever recorded a sandbox. + # Runs without a subject keep NULL either way. A sandbox that + # finishes provisioning afterwards is ignored (status-guarded) + # and released by the dispatcher, so no orphan can attach here. + subject_released_at=case( + (AutomationRun.subject_key.is_not(None), now), + else_=AutomationRun.subject_released_at, + ), ) ) db_result: CursorResult = await session.execute(stmt) # type: ignore[assignment] @@ -1025,15 +1035,11 @@ async def cancel_run( ) # Clean up sandbox for runs that were RUNNING. Cancelling is explicit, so - # unlike `complete_run` the sandbox goes even when the run owns a subject - # -- but the subject is released with it, or the next event would pick this - # run and pay a lookup for a sandbox we just deleted. The key stays on the - # row as the record of what this run was about. + # unlike `complete_run` the sandbox goes even when the run owned a subject + # (already released by the terminal update above, so the next event will + # not pick this run). The key stays on the row as the record of what this + # run was about. if run.sandbox_id: - if run.subject_key and run.subject_released_at is None: - run.subject_released_at = utcnow() - await session.commit() - from openhands.automation.config import get_settings settings = get_settings() diff --git a/openhands/automation/utils/run.py b/openhands/automation/utils/run.py index 8d14cf21..253f9105 100644 --- a/openhands/automation/utils/run.py +++ b/openhands/automation/utils/run.py @@ -293,9 +293,15 @@ async def update_sandbox_id( session_factory: async_sessionmaker[AsyncSession], run_id: uuid.UUID, sandbox_id: str, -) -> None: +) -> bool | None: """Store the sandbox ID on the automation run for later verification. + Only applies while the run is still RUNNING. A run cancelled (or + otherwise finished) mid-provisioning must not gain a sandbox + afterwards: nobody would clean it up, and it would sit on a terminal + row. Returns True when recorded, False when the run is no longer + RUNNING, and None on a write error so execution can continue. + Args: session_factory: Async session factory run_id: The run ID to update @@ -303,14 +309,19 @@ async def update_sandbox_id( """ try: async with session_factory() as session: - await session.execute( + result: CursorResult = await session.execute( # type: ignore[assignment] update(AutomationRun) - .where(AutomationRun.id == run_id) + .where( + AutomationRun.id == run_id, + AutomationRun.status == AutomationRunStatus.RUNNING, + ) .values(sandbox_id=sandbox_id) ) await session.commit() + return (result.rowcount or 0) > 0 except Exception: logger.exception("Failed to update sandbox_id for run %s", run_id) + return None async def update_bash_command_id( diff --git a/tests/test_cancel_run.py b/tests/test_cancel_run.py index ee54032e..de0baa8a 100644 --- a/tests/test_cancel_run.py +++ b/tests/test_cancel_run.py @@ -131,6 +131,114 @@ async def test_cancel_other_orgs_run_returns_403(async_client, async_session): assert resp.status_code == 403 +async def test_cancel_subject_run_without_sandbox_releases_subject( + async_client, async_session +): + """Cancelling a subject run with no sandbox still releases the subject.""" + _, run = await _create_automation_with_run( + async_session, status=AutomationRunStatus.RUNNING + ) + run.subject_key = "team/C123/1755000000.000100" + await async_session.commit() + + resp = await async_client.post(f"/api/automation/v1/runs/{run.id}/cancel") + assert resp.status_code == 200 + + await async_session.refresh(run) + assert run.status == AutomationRunStatus.CANCELLED + assert run.subject_released_at is not None + assert run.subject_key == "team/C123/1755000000.000100" + + +async def test_cancel_ordinary_run_touches_no_subject(async_client, async_session): + """Cancelling a run without a subject leaves subject columns alone.""" + _, run = await _create_automation_with_run( + async_session, status=AutomationRunStatus.RUNNING + ) + + resp = await async_client.post(f"/api/automation/v1/runs/{run.id}/cancel") + assert resp.status_code == 200 + + await async_session.refresh(run) + assert run.status == AutomationRunStatus.CANCELLED + assert run.subject_key is None + assert run.subject_released_at is None + + +async def test_cancelled_subject_no_longer_blocks_resubmission( + async_client, async_session +): + """After cancel, a resubmitted event routes to a new run instead of being + folded into the cancelled run that still holds the subject.""" + from openhands.automation.conversations import continue_conversation + + automation, run = await _create_automation_with_run( + async_session, status=AutomationRunStatus.PENDING + ) + run.subject_key = "team/C123/1755000000.000100" + await async_session.commit() + + resp = await async_client.post(f"/api/automation/v1/runs/{run.id}/cancel") + assert resp.status_code == 200 + + result = await continue_conversation( + async_session, + org_id=TEST_ORG_ID, + source="slack", + subject_key="team/C123/1755000000.000100", + automation_id=automation.id, + event_key="Ev2", + event_payload={}, + ) + assert result.needs_run is True + + +async def test_late_sandbox_record_after_cancel_is_ignored( + async_client, async_session, async_session_factory +): + """A sandbox recorded after cancel must not attach to the cancelled run. + + Covers the race where the dispatcher is still provisioning while the + user cancels: without the status guard the late record would orphan a + sandbox on a terminal row and fork the released subject. + """ + from openhands.automation.conversations import continue_conversation + from openhands.automation.utils.run import update_sandbox_id + + automation, run = await _create_automation_with_run( + async_session, status=AutomationRunStatus.RUNNING + ) + run.subject_key = "team/C123/1755000000.000100" + await async_session.commit() + + resp = await async_client.post(f"/api/automation/v1/runs/{run.id}/cancel") + assert resp.status_code == 200 + + # The endpoint's session stays open in tests (the app commits it on + # teardown in production): close its transaction so the dispatcher's + # own session below sees a committed row instead of blocking on it. + await async_session.commit() + + recorded = await update_sandbox_id(async_session_factory, run.id, "sandbox-A") + assert recorded is False + + await async_session.refresh(run) + assert run.status == AutomationRunStatus.CANCELLED + assert run.sandbox_id is None + assert run.subject_released_at is not None + + result = await continue_conversation( + async_session, + org_id=TEST_ORG_ID, + source="slack", + subject_key="team/C123/1755000000.000100", + automation_id=automation.id, + event_key="Ev2", + event_payload={}, + ) + assert result.needs_run is True + + async def test_cancel_same_org_other_users_run(async_client, async_session): """Cancelling a run owned by another member of the same org should succeed.""" automation = Automation( diff --git a/tests/test_dispatcher.py b/tests/test_dispatcher.py index 20389ed4..8d0bc728 100644 --- a/tests/test_dispatcher.py +++ b/tests/test_dispatcher.py @@ -12,7 +12,10 @@ import pytest from sqlalchemy import select +from sqlalchemy.exc import OperationalError +from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.orm import selectinload +from sqlalchemy.sql.dml import Update from openhands.automation.config import get_config from openhands.automation.conversations import COALESCED_TURNS_KEY @@ -38,6 +41,7 @@ update_run_timeout_at, ) from openhands.automation.utils.tarball_validation import is_http_url +from openhands.automation.utils.time import ensure_utc # Test UUIDs @@ -1318,6 +1322,171 @@ async def test_generic_context_failure_still_marks_failed( assert updated.status_detail["operation"] == "get_execution_context" assert updated.status_detail["transient"] is False + @patch("openhands.automation.dispatcher.execute_in_context", new_callable=AsyncMock) + async def test_cancel_mid_provisioning_releases_sandbox_without_recording( + self, mock_execute, async_session_factory, mock_settings, mock_client + ): + """A run cancelled while provisioning must not gain a sandbox. + + Simulates cancel landing between sandbox creation and the sandbox-id + record: the provisioned sandbox is released instead of being + attached to the terminal row, so it cannot leak or fork the subject. + """ + from sqlalchemy import update + + async with async_session_factory() as session: + automation = Automation( + user_id=TEST_USER_ID, + org_id=TEST_ORG_ID, + name="Test", + trigger={"type": "cron", "schedule": "* * * * *", "timezone": "UTC"}, + tarball_path="https://example.com/code.tar.gz", + entrypoint="uv run main.py", + enabled=True, + ) + session.add(automation) + await session.commit() + + run = AutomationRun( + automation_id=automation.id, + status=AutomationRunStatus.RUNNING, + started_at=utcnow(), + ) + session.add(run) + await session.commit() + run_id = run.id + + async with async_session_factory() as session: + run = ( + ( + await session.execute( + select(AutomationRun) + .options(selectinload(AutomationRun.automation)) + .where(AutomationRun.id == run_id) + ) + ) + .scalars() + .first() + ) + + async def _cancel_mid_flight(*args, **kwargs): + async with async_session_factory() as session: + await session.execute( + update(AutomationRun) + .where(AutomationRun.id == run_id) + .values( + status=AutomationRunStatus.CANCELLED, + completed_at=utcnow(), + subject_released_at=utcnow(), + ) + ) + await session.commit() + return MagicMock(success=True, bash_command_id="cmd-1", error=None) + + mock_execute.side_effect = _cancel_mid_flight + + backend = MagicMock() + backend.is_local_mode = False + ctx = MagicMock( + agent_url="http://agent.test", sandbox_id="sbx-1", session_key="sk-1" + ) + backend.get_execution_context = AsyncMock(return_value=ctx) + backend.build_env_vars = MagicMock(return_value={}) + backend.get_work_dir = MagicMock(return_value="/workspace") + backend.release_context = AsyncMock() + + with patch("openhands.automation.dispatcher.get_backend", return_value=backend): + await _execute_run(run, mock_settings, async_session_factory, mock_client) + + backend.release_context.assert_called_once() + async with async_session_factory() as session: + updated = ( + ( + await session.execute( + select(AutomationRun).where(AutomationRun.id == run_id) + ) + ) + .scalars() + .first() + ) + assert updated.status == AutomationRunStatus.CANCELLED + assert updated.sandbox_id is None + + @pytest.mark.parametrize("failure_at", ["execute", "commit"]) + @patch("openhands.automation.dispatcher.execute_in_context", new_callable=AsyncMock) + async def test_sandbox_record_error_keeps_running_execution( + self, + mock_execute, + async_session_factory, + mock_settings, + mock_client, + caplog, + failure_at, + ): + """A failed sandbox write must not cancel a successful dispatch.""" + run, run_id, _ = await self._make_running_run(async_session_factory) + run.automation.tarball_path = "https://example.com/code.tar.gz" + run.automation.timeout = 600 + mock_execute.return_value = MagicMock( + success=True, bash_command_id="cmd-1", error=None + ) + + backend = MagicMock() + backend.is_local_mode = False + backend.get_execution_context = AsyncMock( + return_value=MagicMock( + agent_url="http://agent.test", sandbox_id="sbx-1", session_key="sk-1" + ) + ) + backend.build_env_vars = MagicMock(return_value={}) + backend.get_work_dir = MagicMock(return_value="/workspace") + backend.release_context = AsyncMock() + + original_execute = AsyncSession.execute + original_commit = AsyncSession.commit + sandbox_session: AsyncSession | None = None + write_error = OperationalError( + "UPDATE automation_runs", {}, RuntimeError("temporary database failure") + ) + + async def fail_sandbox_execute(session, statement, *args, **kwargs): + nonlocal sandbox_session + if isinstance(statement, Update) and "sandbox_id" in ( + statement.compile().params or {} + ): + sandbox_session = session + if failure_at == "execute": + raise write_error + return await original_execute(session, statement, *args, **kwargs) + + async def fail_sandbox_commit(session): + if failure_at == "commit" and session is sandbox_session: + raise write_error + await original_commit(session) + + dispatched_at = utcnow() + with ( + patch("openhands.automation.dispatcher.get_backend", return_value=backend), + patch.object(AsyncSession, "execute", new=fail_sandbox_execute), + patch.object(AsyncSession, "commit", new=fail_sandbox_commit), + ): + await _execute_run(run, mock_settings, async_session_factory, mock_client) + + assert sandbox_session is not None + assert "Failed to update sandbox_id" in caplog.text + backend.release_context.assert_not_awaited() + async with async_session_factory() as session: + updated = await session.get(AutomationRun, run_id) + assert updated is not None + assert updated.status == AutomationRunStatus.RUNNING + assert updated.completed_at is None + assert updated.sandbox_id is None + assert updated.bash_command_id == "cmd-1" + assert updated.timeout_at is not None + assert ensure_utc(updated.timeout_at) >= dispatched_at + timedelta( + seconds=600 + get_config().sandbox.run_timeout_margin + ) + class TestExecuteRunDerivedConversationId: """A subject-owning run creates its conversation under the derived id."""