diff --git a/migrations/versions/027_add_agent_turn_runs.py b/migrations/versions/027_add_agent_turn_runs.py new file mode 100644 index 00000000..ea24b8c6 --- /dev/null +++ b/migrations/versions/027_add_agent_turn_runs.py @@ -0,0 +1,74 @@ +"""Add service-owned agent-turn runs. + +Revision ID: 027 +Revises: 026 +""" + +import sqlalchemy as sa +from alembic import op + + +revision = "027" +down_revision = "026" +branch_labels = None +depends_on = None + + +def upgrade() -> None: + op.drop_index("ix_automation_runs_subject", table_name="automation_runs") + op.add_column( + "automation_runs", sa.Column("subject_source", sa.String(100), nullable=True) + ) + op.add_column( + "automation_runs", sa.Column("conversation_turn", sa.Text(), nullable=True) + ) + op.add_column( + "automation_runs", + sa.Column("conversation_wake_agent", sa.Boolean(), nullable=True), + ) + op.execute( + """ + UPDATE automation_runs + SET subject_source = ( + SELECT json_extract(automations.trigger, '$.source') + FROM automations + WHERE automations.id = automation_runs.automation_id + ) + WHERE subject_key IS NOT NULL + """ + if op.get_context().dialect.name == "sqlite" + else """ + UPDATE automation_runs AS runs + SET subject_source = automations.trigger ->> 'source' + FROM automations + WHERE automations.id = runs.automation_id + AND runs.subject_key IS NOT NULL + """ + ) + op.create_index( + "ix_automation_runs_subject", + "automation_runs", + ["automation_id", "subject_source", "subject_key", "created_at"], + unique=False, + postgresql_where=sa.text( + "subject_key IS NOT NULL AND subject_released_at IS NULL" + ), + sqlite_where=sa.text("subject_key IS NOT NULL AND subject_released_at IS NULL"), + ) + + +def downgrade() -> None: + op.drop_index("ix_automation_runs_subject", table_name="automation_runs") + op.drop_column("automation_runs", "conversation_turn") + op.drop_column("automation_runs", "conversation_wake_agent") + op.drop_column("automation_runs", "subject_source") + op.create_index( + "ix_automation_runs_subject", + "automation_runs", + ["automation_id", "subject_key", "created_at"], + unique=False, + postgresql_where=sa.text( + "subject_key IS NOT NULL AND subject_released_at IS NULL" + ), + sqlite_where=sa.text("subject_key IS NOT NULL AND subject_released_at IS NULL"), + ) diff --git a/openhands/automation/backends/__init__.py b/openhands/automation/backends/__init__.py index afb83a42..d7a15403 100644 --- a/openhands/automation/backends/__init__.py +++ b/openhands/automation/backends/__init__.py @@ -48,7 +48,8 @@ def get_backend(run: AutomationRun) -> ExecutionBackend: if profile_id and not settings.is_local_mode: raise ValueError("Agent profiles require a configured Agent Server") - if settings.is_local_mode and run.execution_mode == "agent": + uses_conversation = bool(run.conversation_turn) or run.execution_mode == "agent" + if settings.is_local_mode and uses_conversation: if profile_id is None: raise ValueError("Agent execution requires an agent profile") return ConversationAgentServerBackend( diff --git a/openhands/automation/backends/conversation.py b/openhands/automation/backends/conversation.py index 4e1240ad..04db423f 100644 --- a/openhands/automation/backends/conversation.py +++ b/openhands/automation/backends/conversation.py @@ -55,6 +55,8 @@ def __init__( @property def conversation_id(self) -> UUID: + if self._run.conversation_id: + return UUID(self._run.conversation_id) automation = self._run.automation source = (automation.trigger or {}).get("source") if self._run.subject_key and source: diff --git a/openhands/automation/conversations.py b/openhands/automation/conversations.py index a07ee192..5c100a39 100644 --- a/openhands/automation/conversations.py +++ b/openhands/automation/conversations.py @@ -152,6 +152,7 @@ def _clean_key(value: str, origin: str) -> str | None: async def _take_subject_lock( session: AsyncSession, automation_id: uuid.UUID, + source: str, subject_key: str, ) -> None: """Serialise every event for one subject, including those finding no run. @@ -170,7 +171,7 @@ async def _take_subject_lock( # Hashed here, not with `hashtextextended`, so the key does not depend on # a server-side hash staying stable across versions. digest = hashlib.blake2b( - f"{automation_id}/{subject_key}".encode(), digest_size=8 + f"{automation_id}/{source}/{subject_key}".encode(), digest_size=8 ).digest() await session.execute( text("SELECT pg_advisory_xact_lock(:key)").bindparams( @@ -182,6 +183,7 @@ async def _take_subject_lock( async def _lock_subject_run( session: AsyncSession, automation_id: uuid.UUID, + source: str, subject_key: str, ) -> AutomationRun | None: """The most recent run holding this subject, started or not, locked. @@ -199,6 +201,7 @@ async def _lock_subject_run( select(AutomationRun) .where( AutomationRun.automation_id == automation_id, + AutomationRun.subject_source == source, AutomationRun.subject_key == subject_key, AutomationRun.subject_released_at.is_(None), ) @@ -258,13 +261,15 @@ async def continue_conversation( The id is known before the first run finishes, so an event arriving mid-run continues that conversation instead of racing a second run. """ - await _take_subject_lock(session, automation_id, subject_key) + await _take_subject_lock(session, automation_id, source, subject_key) - run = await _lock_subject_run(session, automation_id, subject_key) + run = await _lock_subject_run(session, automation_id, source, subject_key) if run is None: return ContinueResult() - conversation_id = conversation_id_for(org_id, automation_id, source, subject_key) + conversation_id = run.conversation_id or conversation_id_for( + org_id, automation_id, source, subject_key + ) turn = compose_turn(source, event_key, event_payload, override=turn_text) if run.started_at is None: diff --git a/openhands/automation/dispatcher.py b/openhands/automation/dispatcher.py index a2ffb4a6..801b0a82 100644 --- a/openhands/automation/dispatcher.py +++ b/openhands/automation/dispatcher.py @@ -22,7 +22,7 @@ from typing import Any import httpx -from sqlalchemy import select, update +from sqlalchemy import or_, select, update from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker from sqlalchemy.orm import selectinload @@ -46,6 +46,7 @@ from openhands.automation.telemetry import capture_automation_event from openhands.automation.utils import log_extra from openhands.automation.utils.api_key import APIKeyError +from openhands.automation.utils.conversation_turn import run_conversation_turn from openhands.automation.utils.kv import create_kv_token from openhands.automation.utils.run import ( disable_automation, @@ -132,26 +133,31 @@ async def _poll_pending_runs( ``org_id``, and tarball config are available for dispatch. """ settings = get_config().service - active_agent_automations = [] + active: list[tuple[uuid.UUID, uuid.UUID, str | None, str]] = [] agent_capacity: int | None = None if settings.is_local_mode: - active_agent_automations = ( + active = list( ( await session.execute( - select(AutomationRun.automation_id).where( - AutomationRun.status == AutomationRunStatus.RUNNING, - AutomationRun.execution_mode == "agent", - ) + select( + AutomationRun.id, + AutomationRun.automation_id, + AutomationRun.conversation_turn, + AutomationRun.execution_mode, + ).where(AutomationRun.status == AutomationRunStatus.RUNNING) ) ) - .scalars() + .tuples() .all() ) + active_conversations = [ + row for row in active if row[2] is not None or row[3] == "agent" + ] agent_capacity = settings.conversation_max_concurrent_runs - len( - active_agent_automations + active_conversations ) if agent_capacity > 0: - batch_size = min(batch_size, 1, agent_capacity) + batch_size = min(batch_size, agent_capacity) select_query = ( select(AutomationRun) @@ -168,11 +174,25 @@ async def _poll_pending_runs( if agent_capacity is not None and agent_capacity <= 0: # Host-side scripts remain runnable while agent conversations occupy # every bounded runtime slot. - select_query = select_query.where(AutomationRun.execution_mode != "agent") - elif active_agent_automations: select_query = select_query.where( - AutomationRun.automation_id.not_in(active_agent_automations) + AutomationRun.execution_mode != "agent", + AutomationRun.conversation_turn.is_(None), ) + if active: + active_script_automations = { + automation_id + for _, automation_id, conversation_turn, execution_mode in active + if conversation_turn is None and execution_mode == "script" + } + if active_script_automations: + # Do not overlap two host scripts for one definition. Subject turns + # are independent conversations and may fan out concurrently. + select_query = select_query.where( + or_( + AutomationRun.conversation_turn.is_not(None), + AutomationRun.automation_id.not_in(active_script_automations), + ) + ) # Apply row locking for PostgreSQL only (SQLite doesn't support it) if not using_sqlite(): @@ -354,6 +374,71 @@ async def _fail( extra=_log_ctx(sandbox_id=ctx.sandbox_id), ) + # Persist the provisioned context before starting work. Programmatic + # follow-up turns may arrive as soon as the conversation starts. + if ctx.runtime_conversation_id is not None: + run.conversation_id = str(ctx.runtime_conversation_id) + async with session_factory() as link_session: + await link_session.execute( + update(AutomationRun) + .where(AutomationRun.id == run.id) + .values(conversation_id=run.conversation_id) + ) + await link_session.commit() + if ctx.sandbox_id: + run.sandbox_id = ctx.sandbox_id + await update_sandbox_id(session_factory, run.id, ctx.sandbox_id) + + if run.conversation_turn is not None: + assert run.conversation_id is not None + await update_run_current_phase(session_factory, run.id, "Agent is working") + try: + await run_conversation_turn( + ctx, + run.conversation_id, + run.conversation_turn, + wake_agent=run.conversation_wake_agent is not False, + timeout=effective_timeout, + ) + except Exception as exc: + run.subject_released_at = utcnow() + async with session_factory() as release_session: + await release_session.execute( + update(AutomationRun) + .where(AutomationRun.id == run.id) + .values(subject_released_at=run.subject_released_at) + ) + await release_session.commit() + await backend.release_context(client, ctx) + await _fail( + "Subject conversation failed", + status_detail=run_status_detail_from_exception( + exc, + phase=RunStatusPhase.EXECUTION, + source="agent_server", + operation="run_subject_turn", + ), + ) + return + try: + await backend.release_context(client, ctx) + except Exception: + # Conversation history is already persisted outside the runtime; + # the runtime registry's TTL cleanup remains the fallback. + logger.exception( + "Failed to release completed subject runtime", + extra=_log_ctx(sandbox_id=ctx.sandbox_id), + ) + await mark_run_terminal(session_factory, run, AutomationRunStatus.COMPLETED) + await capture_automation_event( + "automation_run_completed", + automation=automation, + run=run, + session_factory=session_factory, + properties={"trigger_source": "subject_turn"}, + ) + return + # 3. Build env vars (must be after get_execution_context for cloud mode API key) callback_url = f"{settings.resolved_base_url.rstrip('/')}/v1/runs/{run_id}/complete" env_vars = backend.build_env_vars() diff --git a/openhands/automation/models.py b/openhands/automation/models.py index ce28f49d..685b339d 100644 --- a/openhands/automation/models.py +++ b/openhands/automation/models.py @@ -8,6 +8,7 @@ from sqlalchemy import ( JSON, BigInteger, + Boolean, DateTime, Enum, Float, @@ -226,6 +227,15 @@ class AutomationRun(Base): # sandbox holding the conversation; the conversation id itself is derived. subject_key: Mapped[str | None] = mapped_column(String(500), nullable=True) + # Namespace of the external subject (for example ``github``). Event runs + # inherit it from their trigger; programmatic subject turns supply it. + subject_source: Mapped[str | None] = mapped_column(String(100), nullable=True) + + # A service-owned agent turn. When present, the dispatcher runs this text + # in the provisioned conversation instead of starting the bundle entrypoint. + conversation_turn: Mapped[str | None] = mapped_column(Text, nullable=True) + conversation_wake_agent: Mapped[bool | None] = mapped_column(Boolean, nullable=True) + # When this run stopped being the subject's routing target -- its sandbox # was deleted, or a turn could not reach it. The key itself stays for the # historical record, so lookups filter on this instead of on its absence. @@ -293,6 +303,7 @@ class AutomationRun(Base): Index( "ix_automation_runs_subject", "automation_id", + "subject_source", "subject_key", "created_at", postgresql_where=(subject_key.isnot(None)) diff --git a/openhands/automation/utils/conversation_turn.py b/openhands/automation/utils/conversation_turn.py index 0e933edf..abc6fa24 100644 --- a/openhands/automation/utils/conversation_turn.py +++ b/openhands/automation/utils/conversation_turn.py @@ -21,6 +21,7 @@ import httpx from openhands.automation.backends import get_backend +from openhands.automation.backends.base import ExecutionContext from openhands.automation.backends.local import LocalAgentServerBackend from openhands.automation.config import get_config from openhands.automation.models import AutomationRun @@ -250,6 +251,54 @@ def _send_turn( workspace.reset_client() +def _run_turn( + context: ExecutionContext, + conversation_id: str, + text: str, + wake_agent: bool, + timeout: int, +) -> None: + workspace = RemoteWorkspace( + host=context.agent_url, + api_key=context.session_key, + working_dir="/", + runtime_conversation_id=context.runtime_conversation_id, + ) + conversation = None + try: + conversation = RemoteConversation.attach( + workspace=workspace, + conversation_id=UUID(conversation_id), + visualizer=None, + ) + conversation.send_message(text) + if wake_agent: + conversation.run(timeout=timeout) + finally: + if conversation is not None: + conversation.close() + workspace.reset_client() + + +async def run_conversation_turn( + context: ExecutionContext, + conversation_id: str, + text: str, + *, + wake_agent: bool, + timeout: int, +) -> None: + """Run one service-owned turn in an already provisioned conversation.""" + await asyncio.to_thread( + _run_turn, + context, + conversation_id, + text, + wake_agent, + timeout, + ) + + async def send_conversation_turn( run: AutomationRun, conversation_id: str, diff --git a/openhands/automation/utils/run.py b/openhands/automation/utils/run.py index edd96148..dfc29eb8 100644 --- a/openhands/automation/utils/run.py +++ b/openhands/automation/utils/run.py @@ -15,6 +15,7 @@ AutomationRun, AutomationRunStatus, ) +from openhands.automation.subjects import conversation_id_for from openhands.automation.telemetry import capture_automation_event from openhands.automation.utils.time import utcnow from openhands.automation.utils.timeout import resolve_automation_timeout_seconds @@ -30,6 +31,35 @@ ) +def create_agent_turn_run( + requester: AutomationRun, + *, + source: str, + subject_key: str, + turn: str, + wake_agent: bool, +) -> AutomationRun: + """Build a tracked agent run selected by an ordinary automation script.""" + if requester.agent_profile_id is None: + raise ValueError("Agent turns require an agent profile") + automation = requester.automation + return AutomationRun( + id=uuid.uuid4(), + automation_id=requester.automation_id, + agent_profile_id=requester.agent_profile_id, + execution_mode="agent", + telemetry_distinct_id=requester.telemetry_distinct_id, + status=AutomationRunStatus.PENDING, + subject_source=source, + subject_key=subject_key, + conversation_id=conversation_id_for( + automation.org_id, automation.id, source, subject_key + ), + conversation_turn=turn, + conversation_wake_agent=wake_agent, + ) + + async def disable_automation( session_factory: async_sessionmaker[AsyncSession], automation_id: uuid.UUID, diff --git a/openhands/automation/utils/webhook.py b/openhands/automation/utils/webhook.py index 0b800f6e..45ff41c2 100644 --- a/openhands/automation/utils/webhook.py +++ b/openhands/automation/utils/webhook.py @@ -268,6 +268,9 @@ async def create_automation_run( telemetry_distinct_id=automation.telemetry_distinct_id, agent_profile_id=automation.agent_profile_id, subject_key=subject_key, + subject_source=( + (automation.trigger or {}).get("source") if subject_key else None + ), ) session.add(run) return run diff --git a/tests/test_backends.py b/tests/test_backends.py index d5fac937..502a6210 100644 --- a/tests/test_backends.py +++ b/tests/test_backends.py @@ -456,6 +456,7 @@ def mock_run(self): run = MagicMock() run.sandbox_id = "sandbox-123" run.agent_profile_id = None + run.conversation_turn = None run.execution_mode = "script" return run diff --git a/tests/test_conversation_turn.py b/tests/test_conversation_turn.py index 9d287840..da1a6e20 100644 --- a/tests/test_conversation_turn.py +++ b/tests/test_conversation_turn.py @@ -12,11 +12,13 @@ import httpx import pytest +from openhands.automation.backends.base import ExecutionContext from openhands.automation.backends.conversation import ConversationAgentServerBackend from openhands.automation.backends.local import LocalAgentServerBackend from openhands.automation.models import AutomationRun from openhands.automation.utils import conversation_turn as turn_module, utcnow from openhands.automation.utils.conversation_turn import ( + _run_turn, compose_turn, send_conversation_turn, ) @@ -25,6 +27,36 @@ CONVERSATION_ID = str(uuid4()) +def test_service_owned_turn_keeps_the_runtime_scope(monkeypatch): + """Agent tools stay in the conversation's selected local/Docker runtime.""" + from unittest.mock import MagicMock + + runtime_id = uuid4() + workspace = MagicMock() + remote_workspace = MagicMock(return_value=workspace) + conversation = MagicMock() + monkeypatch.setattr(turn_module, "RemoteWorkspace", remote_workspace) + monkeypatch.setattr( + turn_module.RemoteConversation, "attach", MagicMock(return_value=conversation) + ) + + _run_turn( + ExecutionContext( + agent_url="https://agent.example.com", + session_key="outer-key", + runtime_conversation_id=runtime_id, + ), + CONVERSATION_ID, + "work on this", + True, + 120, + ) + + assert remote_workspace.call_args.kwargs["runtime_conversation_id"] == runtime_id + conversation.send_message.assert_called_once_with("work on this") + conversation.run.assert_called_once_with(timeout=120) + + @pytest.fixture def conversation_transport(sdk_http_transport, monkeypatch): from unittest.mock import MagicMock diff --git a/tests/test_conversations.py b/tests/test_conversations.py index c6b870f5..12f27dd1 100644 --- a/tests/test_conversations.py +++ b/tests/test_conversations.py @@ -1098,6 +1098,7 @@ def handler(request: httpx.Request) -> httpx.Response: status=AutomationRunStatus.COMPLETED, started_at=utcnow(), sandbox_id="sbx-1", + subject_source="slack", subject_key=subject_key, ) ) @@ -1295,7 +1296,7 @@ async def test_the_subject_lock_orders_events_that_find_no_run( async def hold(name: str, work: float) -> None: async with async_session_factory() as session: - await _take_subject_lock(session, automation_id, "T1/C1/1.1") + await _take_subject_lock(session, automation_id, "slack", "T1/C1/1.1") order.append(f"{name} in") await asyncio.sleep(work) order.append(f"{name} out") @@ -1331,7 +1332,7 @@ async def test_a_different_subject_is_not_blocked_by_a_held_lock( async def hold(name: str, subject_key: str, work: float) -> None: async with async_session_factory() as session: - await _take_subject_lock(session, automation_id, subject_key) + await _take_subject_lock(session, automation_id, "slack", subject_key) order.append(f"{name} in") await asyncio.sleep(work) order.append(f"{name} out") diff --git a/tests/test_db.py b/tests/test_db.py index ebef9266..552cd281 100644 --- a/tests/test_db.py +++ b/tests/test_db.py @@ -261,6 +261,9 @@ def test_migrations_run_on_sqlite(self, monkeypatch): assert "cost" in run_columns assert "agent_profile_id" in run_columns assert "execution_mode" in run_columns + assert "subject_source" in run_columns + assert "conversation_turn" in run_columns + assert "conversation_wake_agent" in run_columns assert "agent_profile_id" in { column["name"] for column in inspector.get_columns("automations") } diff --git a/tests/test_dispatcher.py b/tests/test_dispatcher.py index ca55f355..219575a2 100644 --- a/tests/test_dispatcher.py +++ b/tests/test_dispatcher.py @@ -29,6 +29,7 @@ from openhands.automation.subjects import conversation_id_for from openhands.automation.utils import utcnow from openhands.automation.utils.run import ( + create_agent_turn_run, mark_run_status, mark_run_terminal, update_run_current_phase, @@ -75,6 +76,55 @@ def test_gs_url_is_not_http(self): assert is_http_url("gs://bucket/key.tar.gz") is False +@pytest.mark.asyncio +async def test_local_poller_fans_out_subject_agents_from_one_automation( + async_session, monkeypatch +): + from openhands.automation.config import clear_config_cache + + monkeypatch.setenv("AUTOMATION_AGENT_SERVER_URL", "http://agent.test") + monkeypatch.setenv("AUTOMATION_CONVERSATION_MAX_CONCURRENT_RUNS", "4") + clear_config_cache() + try: + automation = Automation( + user_id=TEST_USER_ID, + org_id=TEST_ORG_ID, + name="Developer scanner", + trigger={"type": "cron", "schedule": "* * * * *", "timezone": "UTC"}, + tarball_path="https://example.com/scanner.tar.gz", + entrypoint="python scanner.py", + ) + async_session.add(automation) + await async_session.flush() + async_session.add( + AutomationRun( + automation_id=automation.id, + status=AutomationRunStatus.RUNNING, + conversation_turn="Implement issue 1", + ) + ) + for number in (2, 3, 4): + async_session.add( + AutomationRun( + automation_id=automation.id, + status=AutomationRunStatus.PENDING, + conversation_turn=f"Implement issue {number}", + ) + ) + await async_session.commit() + + pending = await _poll_pending_runs(async_session, batch_size=10) + + assert {run.conversation_turn for run in pending} == { + "Implement issue 2", + "Implement issue 3", + "Implement issue 4", + } + finally: + clear_config_cache() + + +@pytest.mark.asyncio async def test_script_dispatch_is_not_blocked_by_full_agent_capacity( async_session_factory, monkeypatch ): @@ -134,6 +184,70 @@ async def test_script_dispatch_is_not_blocked_by_full_agent_capacity( clear_config_cache() +@pytest.mark.asyncio +async def test_agent_turn_run_skips_bundle_and_releases_runtime( + async_session_factory, mock_settings, mock_client +): + async with async_session_factory() as session: + automation = Automation( + user_id=TEST_USER_ID, + org_id=TEST_ORG_ID, + name="Developer scanner", + trigger={"type": "cron", "schedule": "* * * * *", "timezone": "UTC"}, + tarball_path="https://example.com/scanner.tar.gz", + entrypoint="python scanner.py", + ) + requester = AutomationRun( + automation=automation, + status=AutomationRunStatus.RUNNING, + agent_profile_id=uuid.uuid4(), + ) + session.add(requester) + await session.flush() + run = create_agent_turn_run( + requester, + source="github", + subject_key="repository-42/issue-9", + turn="Implement issue 9", + wake_agent=True, + ) + run.status = AutomationRunStatus.RUNNING + run.started_at = utcnow() + assert run.conversation_id is not None + conversation_id = uuid.UUID(run.conversation_id) + session.add(run) + await session.commit() + run_id = run.id + await session.refresh(run, attribute_names=["automation"]) + + backend = AsyncMock() + backend.get_execution_context.return_value = ExecutionContext( + agent_url="http://agent.test", + session_key="runtime-key", + runtime_conversation_id=conversation_id, + ) + with ( + patch("openhands.automation.dispatcher.get_backend", return_value=backend), + patch( + "openhands.automation.dispatcher.run_conversation_turn", + new_callable=AsyncMock, + ) as run_turn, + patch( + "openhands.automation.dispatcher.execute_in_context", + new_callable=AsyncMock, + ) as execute_bundle, + ): + await _execute_run(run, mock_settings, async_session_factory, mock_client) + + run_turn.assert_awaited_once() + backend.release_context.assert_awaited_once() + execute_bundle.assert_not_awaited() + async with async_session_factory() as session: + finished = await session.get(AutomationRun, run_id) + assert finished is not None + assert finished.status == AutomationRunStatus.COMPLETED + + class TestMarkRunStatus: """Tests for mark_run_status function."""