From 666ade6efaad24e93c5fbaf9ca13ecb2ba4d7d7b Mon Sep 17 00:00:00 2001 From: openhands Date: Wed, 24 Jun 2026 21:15:10 +0000 Subject: [PATCH] Add post-run automation callbacks Co-authored-by: openhands --- migrations/versions/010_post_run_callbacks.py | 106 +++++ openhands/automation/backends/base.py | 11 + openhands/automation/backends/cloud.py | 22 ++ openhands/automation/backends/local.py | 7 + openhands/automation/callbacks.py | 368 ++++++++++++++++++ openhands/automation/models.py | 72 ++++ openhands/automation/preset_router.py | 62 ++- openhands/automation/router.py | 91 ++++- openhands/automation/schemas.py | 127 +++++- openhands/automation/watchdog.py | 60 ++- tests/test_callbacks.py | 291 ++++++++++++++ tests/test_preset_router.py | 39 ++ tests/test_schemas.py | 49 +++ 13 files changed, 1278 insertions(+), 27 deletions(-) create mode 100644 migrations/versions/010_post_run_callbacks.py create mode 100644 openhands/automation/callbacks.py create mode 100644 tests/test_callbacks.py diff --git a/migrations/versions/010_post_run_callbacks.py b/migrations/versions/010_post_run_callbacks.py new file mode 100644 index 00000000..6b3f2005 --- /dev/null +++ b/migrations/versions/010_post_run_callbacks.py @@ -0,0 +1,106 @@ +"""Add post-run callback configuration and execution records. + +Revision ID: 010 +Revises: 009 +Create Date: 2026-06-24 +""" + +from collections.abc import Sequence + +import sqlalchemy as sa +from alembic import op + + +revision: str = "010" +down_revision: str = "009" +branch_labels: str | Sequence[str] | None = None +depends_on: str | Sequence[str] | None = None + + +def _is_sqlite() -> bool: + return op.get_bind().dialect.name == "sqlite" + + +def upgrade() -> None: + op.add_column("automations", sa.Column("callbacks", sa.JSON(), nullable=True)) + + op.create_table( + "automation_run_callbacks", + sa.Column("id", sa.Uuid(), nullable=False), + sa.Column("run_id", sa.Uuid(), nullable=False), + sa.Column("name", sa.String(length=100), nullable=False), + sa.Column("trigger_status", sa.String(length=20), nullable=False), + sa.Column("entrypoint", sa.Text(), nullable=False), + sa.Column("timeout", sa.Integer(), nullable=True), + sa.Column("status", sa.String(length=20), nullable=False), + sa.Column("bash_command_id", sa.String(length=64), nullable=True), + sa.Column("error_detail", sa.Text(), nullable=True), + sa.Column("order", sa.Integer(), nullable=False), + sa.Column("started_at", sa.DateTime(timezone=True), nullable=True), + sa.Column("timeout_at", sa.DateTime(timezone=True), nullable=True), + sa.Column("completed_at", sa.DateTime(timezone=True), nullable=True), + sa.Column( + "created_at", + sa.DateTime(timezone=True), + server_default=sa.text("CURRENT_TIMESTAMP"), + nullable=False, + ), + sa.ForeignKeyConstraint(["run_id"], ["automation_runs.id"], ondelete="CASCADE"), + sa.PrimaryKeyConstraint("id"), + ) + op.create_index( + "ix_automation_run_callbacks_run_id", + "automation_run_callbacks", + ["run_id"], + ) + op.create_index( + "ix_automation_run_callbacks_status", + "automation_run_callbacks", + ["status"], + ) + op.create_index( + "ix_automation_run_callbacks_timeout_at", + "automation_run_callbacks", + ["timeout_at"], + ) + op.create_index( + "ix_automation_run_callbacks_run_order", + "automation_run_callbacks", + ["run_id", "order"], + ) + op.create_index( + "ix_automation_run_callbacks_status_timeout", + "automation_run_callbacks", + ["status", "timeout_at"], + ) + + if not _is_sqlite(): + op.execute( + "COMMENT ON COLUMN automations.callbacks IS " + "'Post-run callback configuration for automation runs.'" + ) + + +def downgrade() -> None: + op.drop_index( + "ix_automation_run_callbacks_status_timeout", + table_name="automation_run_callbacks", + ) + op.drop_index( + "ix_automation_run_callbacks_run_order", + table_name="automation_run_callbacks", + ) + op.drop_index( + "ix_automation_run_callbacks_timeout_at", + table_name="automation_run_callbacks", + ) + op.drop_index( + "ix_automation_run_callbacks_status", + table_name="automation_run_callbacks", + ) + op.drop_index( + "ix_automation_run_callbacks_run_id", + table_name="automation_run_callbacks", + ) + op.drop_table("automation_run_callbacks") + op.drop_column("automations", "callbacks") diff --git a/openhands/automation/backends/base.py b/openhands/automation/backends/base.py index 166985a9..fb1a7161 100644 --- a/openhands/automation/backends/base.py +++ b/openhands/automation/backends/base.py @@ -64,6 +64,17 @@ async def get_execution_context( TimeoutError: If sandbox doesn't become ready in time (Cloud mode) """ + @abstractmethod + async def get_existing_execution_context( + self, client: httpx.AsyncClient + ) -> ExecutionContext: + """Reconnect to an existing run execution context. + + Used by post-run callbacks after the main run has already dispatched. + For Cloud mode this discovers the still-running sandbox by run.sandbox_id. + For Local mode this returns the configured persistent agent server. + """ + @abstractmethod async def release_context( self, client: httpx.AsyncClient, ctx: ExecutionContext diff --git a/openhands/automation/backends/cloud.py b/openhands/automation/backends/cloud.py index 22567d85..ae784c05 100644 --- a/openhands/automation/backends/cloud.py +++ b/openhands/automation/backends/cloud.py @@ -27,6 +27,7 @@ from openhands.automation.utils.sandbox import ( cleanup_sandbox, delete_sandbox, + get_sandbox_agent_url, verify_run_status, ) @@ -174,6 +175,27 @@ async def _do_acquire() -> tuple[str, str, str]: api_key=await self._ensure_api_key(), ) + async def get_existing_execution_context( + self, client: httpx.AsyncClient + ) -> ExecutionContext: + """Discover the existing sandbox's agent server context.""" + sandbox_id = self._run.sandbox_id + if not sandbox_id: + raise RuntimeError("Run has no sandbox_id for callback execution") + + api_key = await self._ensure_api_key() + result = await get_sandbox_agent_url(client, self.api_url, api_key, sandbox_id) + if result is None: + raise RuntimeError(f"Sandbox {sandbox_id} is not available") + agent_url, session_key = result + return ExecutionContext( + agent_url=agent_url, + session_key=session_key, + sandbox_id=sandbox_id, + api_url=self.api_url, + api_key=api_key, + ) + async def release_context( self, client: httpx.AsyncClient, ctx: ExecutionContext ) -> None: diff --git a/openhands/automation/backends/local.py b/openhands/automation/backends/local.py index b6909cd6..ab01265c 100644 --- a/openhands/automation/backends/local.py +++ b/openhands/automation/backends/local.py @@ -112,6 +112,13 @@ async def get_execution_context( sandbox_id=None, # No sandbox in local mode ) + async def get_existing_execution_context( + self, + client: httpx.AsyncClient, # noqa: ARG002 + ) -> ExecutionContext: + """Return the persistent local agent server context.""" + return await self.get_execution_context(client) + async def release_context( self, client: httpx.AsyncClient, # noqa: ARG002 diff --git a/openhands/automation/callbacks.py b/openhands/automation/callbacks.py new file mode 100644 index 00000000..0c970130 --- /dev/null +++ b/openhands/automation/callbacks.py @@ -0,0 +1,368 @@ +"""Post-run callback scheduling and dispatch.""" + +from __future__ import annotations + +import json +import logging +import uuid +from datetime import timedelta +from typing import Any + +import httpx +from sqlalchemy import func, select, update +from sqlalchemy.ext.asyncio import AsyncSession + +from openhands.automation.backends import get_backend +from openhands.automation.config import get_config +from openhands.automation.execution import _shell_quote, _start_bash +from openhands.automation.models import ( + Automation, + AutomationRun, + AutomationRunCallback, + AutomationRunCallbackStatus, + AutomationRunStatus, +) +from openhands.automation.schemas import CallbackCompleteItem +from openhands.automation.utils import log_extra, utcnow + + +logger = logging.getLogger("automation.callbacks") + +_TERMINAL_CALLBACK_STATUSES = { + AutomationRunCallbackStatus.COMPLETED, + AutomationRunCallbackStatus.FAILED, + AutomationRunCallbackStatus.SKIPPED, +} + + +def _matching_callbacks( + automation: Automation, terminal_status: AutomationRunStatus +) -> list[tuple[int, dict[str, Any]]]: + status_value = terminal_status.value + return [ + (idx, config) + for idx, config in enumerate(automation.callbacks or []) + if status_value in set(config.get("on") or []) + ] + + +def _callback_command_from_config(config: dict[str, Any]) -> str: + entrypoint = config.get("entrypoint") + if entrypoint: + return str(entrypoint) + + inline_python = config.get("inline_python") + if inline_python: + marker = f"AUTOMATION_CALLBACK_{uuid.uuid4().hex}" + return f"python - <<'{marker}'\n{inline_python}\n{marker}" + + raise ValueError("callback config must include entrypoint or inline_python") + + +def _callback_complete_url(run_id: uuid.UUID) -> str: + base_url = get_config().service.resolved_base_url.rstrip("/") + return f"{base_url}/v1/runs/{run_id}/callbacks/complete" + + +def _chain_timeout(callbacks: list[AutomationRunCallback]) -> int: + sandbox_timeout = get_config().sandbox.max_run_duration + total = 0 + for callback in callbacks: + total += callback.timeout or sandbox_timeout + return max(total, 1) + + +def _build_callback_chain_command( + *, + work_dir: str, + run: AutomationRun, + callbacks: list[AutomationRunCallback], +) -> str: + callback_specs = [ + { + "id": str(callback.id), + "name": callback.name, + "entrypoint": callback.entrypoint, + "timeout": callback.timeout, + } + for callback in callbacks + ] + payload = { + "work_dir": work_dir, + "callback_url": _callback_complete_url(run.id), + "callbacks": callback_specs, + "base_env": { + "AUTOMATION_RUN_ID": str(run.id), + "AUTOMATION_MAIN_STATUS": run.status.value, + "AUTOMATION_MAIN_ERROR": run.error_detail or "", + "AUTOMATION_CONVERSATION_ID": run.conversation_id or "", + "AUTOMATION_EVENT_PAYLOAD": json.dumps(run.event_payload or {}), + }, + } + marker = f"AUTOMATION_CALLBACK_CHAIN_{uuid.uuid4().hex}" + script = f""" +import json +import os +import subprocess +import urllib.request +from datetime import datetime, timezone + +config = {json.dumps(payload)!r} +config = json.loads(config) +results = [] +base_env = os.environ.copy() +base_env.update(config["base_env"]) + +for callback in config["callbacks"]: + env = base_env.copy() + env["AUTOMATION_CALLBACK_NAME"] = callback["name"] + started_at = datetime.now(timezone.utc).isoformat() + exit_code = None + stdout = "" + stderr = "" + error_detail = None + try: + completed = subprocess.run( + callback["entrypoint"], + shell=True, + cwd=config["work_dir"], + env=env, + capture_output=True, + text=True, + timeout=callback.get("timeout"), + ) + exit_code = completed.returncode + stdout = completed.stdout or "" + stderr = completed.stderr or "" + if exit_code != 0: + parts = [f"exit_code={{exit_code}}"] + if stderr: + parts.append(f"stderr: {{stderr[-1000:]}}") + if stdout: + parts.append(f"stdout: {{stdout[-500:]}}") + error_detail = "\\n".join(parts) + except subprocess.TimeoutExpired as exc: + stdout = (exc.stdout or "") if isinstance(exc.stdout, str) else "" + stderr = (exc.stderr or "") if isinstance(exc.stderr, str) else "" + error_detail = "Timed out" + except Exception as exc: + error_detail = str(exc) + + results.append({{ + "id": callback["id"], + "name": callback["name"], + "status": "COMPLETED" if exit_code == 0 and error_detail is None else "FAILED", + "exit_code": exit_code, + "stdout": stdout[-500:], + "stderr": stderr[-1000:], + "error_detail": error_detail, + "started_at": started_at, + "completed_at": datetime.now(timezone.utc).isoformat(), + }}) + +body = json.dumps({{"callbacks": results}}).encode("utf-8") +token = ( + os.environ.get("AUTOMATION_CALLBACK_API_KEY") + or os.environ.get("OPENHANDS_API_KEY") + or "" +) +headers = {{"Content-Type": "application/json"}} +if token: + headers["Authorization"] = f"Bearer {{token}}" +request = urllib.request.Request( + config["callback_url"], + data=body, + headers=headers, + method="POST", +) +try: + with urllib.request.urlopen(request, timeout=30) as response: + response.read() +except Exception as exc: + print(f"Failed to report callback completion: {{exc}}", flush=True) + raise +""".strip() + return f"cd {_shell_quote(work_dir)} && python - <<'{marker}'\n{script}\n{marker}" + + +async def _get_automation( + session: AsyncSession, run: AutomationRun +) -> Automation | None: + if run.automation is not None: + return run.automation + return await session.get(Automation, run.automation_id) + + +async def schedule_and_dispatch_callbacks_for_run( + session: AsyncSession, + run: AutomationRun, +) -> int: + """Create callback records and start one in-sandbox callback chain. + + The records are committed before the in-sandbox wrapper starts so a fast + callback completion request can update durable rows immediately. + """ + if run.status not in (AutomationRunStatus.COMPLETED, AutomationRunStatus.FAILED): + return 0 + + automation = await _get_automation(session, run) + if automation is None: + return 0 + + matches = _matching_callbacks(automation, run.status) + if not matches: + return 0 + + existing = await session.scalar( + select(func.count()) + .select_from(AutomationRunCallback) + .where(AutomationRunCallback.run_id == run.id) + ) + if existing: + return 0 + + records: list[AutomationRunCallback] = [] + now = utcnow() + for order, config in matches: + try: + entrypoint = _callback_command_from_config(config) + status = AutomationRunCallbackStatus.PENDING + error_detail = None + except ValueError as exc: + entrypoint = "" + status = AutomationRunCallbackStatus.SKIPPED + error_detail = str(exc) + + if not run.sandbox_id and not get_config().service.is_local_mode: + status = AutomationRunCallbackStatus.SKIPPED + error_detail = "Run has no sandbox for callback execution" + + record = AutomationRunCallback( + run_id=run.id, + name=str(config.get("name") or f"callback-{order}"), + trigger_status=run.status, + entrypoint=entrypoint, + timeout=config.get("timeout"), + status=status, + error_detail=error_detail, + order=order, + completed_at=now if status == AutomationRunCallbackStatus.SKIPPED else None, + ) + records.append(record) + session.add(record) + + await session.flush() + runnable = [r for r in records if r.status == AutomationRunCallbackStatus.PENDING] + await session.commit() + + if not runnable: + return len(records) + + backend = get_backend(run) + extra = log_extra(run_id=str(run.id), sandbox_id=run.sandbox_id) + try: + async with httpx.AsyncClient(timeout=60.0) as client: + ctx = await backend.get_existing_execution_context(client) + command = _build_callback_chain_command( + work_dir=backend.get_work_dir(str(run.id)), + run=run, + callbacks=runnable, + ) + chain_timeout = _chain_timeout(runnable) + command_id = await _start_bash( + client, + ctx.agent_url, + ctx.session_key, + command, + timeout=chain_timeout, + ) + + started_at = utcnow() + await session.execute( + update(AutomationRunCallback) + .where(AutomationRunCallback.id.in_([r.id for r in runnable])) + .values( + status=AutomationRunCallbackStatus.RUNNING, + bash_command_id=command_id, + started_at=started_at, + timeout_at=started_at + timedelta(seconds=chain_timeout), + ) + ) + await session.commit() + logger.info( + "Callback chain started (command_id=%s)", + command_id, + extra=extra, + ) + except Exception as exc: + logger.warning("Callback chain failed to start: %s", exc, extra=extra) + failed_at = utcnow() + await session.execute( + update(AutomationRunCallback) + .where(AutomationRunCallback.id.in_([r.id for r in runnable])) + .values( + status=AutomationRunCallbackStatus.FAILED, + completed_at=failed_at, + error_detail=str(exc), + ) + ) + await session.commit() + + return len(records) + + +async def complete_callback_records( + session: AsyncSession, + run: AutomationRun, + callbacks: list[CallbackCompleteItem], +) -> None: + """Persist completion results reported by the in-sandbox callback wrapper.""" + for result in callbacks: + status = AutomationRunCallbackStatus(result.status) + await session.execute( + update(AutomationRunCallback) + .where( + AutomationRunCallback.id == result.id, + AutomationRunCallback.run_id == run.id, + ) + .values( + status=status, + completed_at=result.completed_at or utcnow(), + error_detail=result.error_detail, + ) + ) + + +async def run_has_unfinished_callbacks(session: AsyncSession, run_id: Any) -> bool: + count = await session.scalar( + select(func.count()) + .select_from(AutomationRunCallback) + .where( + AutomationRunCallback.run_id == run_id, + AutomationRunCallback.status.notin_(list(_TERMINAL_CALLBACK_STATUSES)), + ) + ) + return bool(count) + + +async def cleanup_run_after_callbacks_if_ready( + session: AsyncSession, + run: AutomationRun, +) -> bool: + """Clean up explicit-cleanup runs once all callbacks are terminal.""" + automation = await _get_automation(session, run) + if automation is None or automation.keep_alive is True or not run.sandbox_id: + return False + if await run_has_unfinished_callbacks(session, run.id): + return False + + try: + await get_backend(run).cleanup_after_verification(str(run.id)) + return True + except Exception as exc: + logger.warning( + "Cleanup after callbacks failed: %s", + exc, + extra=log_extra(run_id=str(run.id), sandbox_id=run.sandbox_id), + ) + return False diff --git a/openhands/automation/models.py b/openhands/automation/models.py index 7ffb0999..ffbdd652 100644 --- a/openhands/automation/models.py +++ b/openhands/automation/models.py @@ -44,6 +44,16 @@ class AutomationRunStatus(enum.Enum): SKIPPED = "SKIPPED" +class AutomationRunCallbackStatus(enum.Enum): + """Status of a post-run callback execution.""" + + PENDING = "PENDING" + RUNNING = "RUNNING" + COMPLETED = "COMPLETED" + FAILED = "FAILED" + SKIPPED = "SKIPPED" + + class Automation(Base): """An automation definition: what to run and when to trigger it.""" @@ -82,6 +92,10 @@ class Automation(Base): # means the automation service owns explicit cleanup. keep_alive: Mapped[bool | None] = mapped_column(default=None, nullable=True) + # Post-run callback configuration. Each item contains name, on, entrypoint, + # timeout, and optionally inline_python before preset materialization. + callbacks: Mapped[list[dict] | None] = mapped_column(JSON, nullable=True) + # Whether the automation is enabled (can be triggered) enabled: Mapped[bool] = mapped_column(default=True, nullable=False, index=True) @@ -185,6 +199,9 @@ class AutomationRun(Base): # Relationship back to automation automation: Mapped["Automation"] = relationship("Automation", back_populates="runs") + callbacks: Mapped[list["AutomationRunCallback"]] = relationship( + "AutomationRunCallback", back_populates="run", cascade="all, delete-orphan" + ) __table_args__ = ( # Partial index for efficient PENDING polling. @@ -200,6 +217,61 @@ class AutomationRun(Base): ) +class AutomationRunCallback(Base): + """A post-run callback execution for an automation run.""" + + __tablename__ = "automation_run_callbacks" + + id: Mapped[uuid.UUID] = mapped_column(Uuid, primary_key=True, default=uuid.uuid4) + run_id: Mapped[uuid.UUID] = mapped_column( + Uuid, + ForeignKey("automation_runs.id", ondelete="CASCADE"), + nullable=False, + index=True, + ) + + name: Mapped[str] = mapped_column(String(100), nullable=False) + trigger_status: Mapped[AutomationRunStatus] = mapped_column( + Enum(AutomationRunStatus, native_enum=False, length=20), + nullable=False, + ) + entrypoint: Mapped[str] = mapped_column(Text, nullable=False) + timeout: Mapped[int | None] = mapped_column(nullable=True) + status: Mapped[AutomationRunCallbackStatus] = mapped_column( + Enum(AutomationRunCallbackStatus, native_enum=False, length=20), + nullable=False, + default=AutomationRunCallbackStatus.PENDING, + index=True, + ) + bash_command_id: Mapped[str | None] = mapped_column(String(64), nullable=True) + error_detail: Mapped[str | None] = mapped_column(Text, nullable=True) + order: Mapped[int] = mapped_column(nullable=False, default=0) + + started_at: Mapped[datetime | None] = mapped_column( + DateTime(timezone=True), nullable=True + ) + timeout_at: Mapped[datetime | None] = mapped_column( + DateTime(timezone=True), nullable=True, index=True + ) + completed_at: Mapped[datetime | None] = mapped_column( + DateTime(timezone=True), nullable=True + ) + created_at: Mapped[datetime] = mapped_column( + DateTime(timezone=True), + server_default=text("CURRENT_TIMESTAMP"), + nullable=False, + ) + + run: Mapped["AutomationRun"] = relationship( + "AutomationRun", back_populates="callbacks" + ) + + __table_args__ = ( + Index("ix_automation_run_callbacks_run_order", "run_id", "order"), + Index("ix_automation_run_callbacks_status_timeout", "status", "timeout_at"), + ) + + class TarballUpload(Base): """A tarball upload for automation code. diff --git a/openhands/automation/preset_router.py b/openhands/automation/preset_router.py index a95a9fab..523f643a 100644 --- a/openhands/automation/preset_router.py +++ b/openhands/automation/preset_router.py @@ -28,7 +28,7 @@ from openhands.automation.constants import MODEL_PROFILE_PATTERN from openhands.automation.db import get_session from openhands.automation.models import Automation, TarballUpload, UploadStatus -from openhands.automation.schemas import AutomationResponse, Trigger +from openhands.automation.schemas import AutomationCallback, AutomationResponse, Trigger from openhands.automation.storage import FileStore, get_file_store from openhands.automation.utils import utcnow from openhands.automation.utils.model_profiles import resolve_model_profile_for_user @@ -146,6 +146,10 @@ class CreatePromptAutomationRequest(BaseModel): "completion (or after post-run callbacks, when configured)." ), ) + callbacks: list[AutomationCallback] | None = Field( + default=None, + description="Post-run callbacks to run after COMPLETED or FAILED runs.", + ) repos: list[RepoSource] | None = Field( default=None, description=( @@ -177,7 +181,31 @@ def _add_file_to_tar( tar.addfile(info, io.BytesIO(content_bytes)) -def _generate_tarball(prompt: str, repos: list[RepoSource] | None = None) -> bytes: +def _materialize_callback_files( + callbacks: list[AutomationCallback] | None, +) -> tuple[list[dict[str, Any]] | None, dict[str, str]]: + """Return persisted callback config and files for inline preset callbacks.""" + if not callbacks: + return None, {} + + persisted: list[dict[str, Any]] = [] + files: dict[str, str] = {} + for callback in callbacks: + data = callback.model_dump(exclude_none=True) + inline_python = data.pop("inline_python", None) + if inline_python is not None: + path = f"callbacks/{callback.name}.py" + files[path] = inline_python + data["entrypoint"] = f"python {path}" + persisted.append(data) + return persisted, files + + +def _generate_tarball( + prompt: str, + repos: list[RepoSource] | None = None, + callback_files: dict[str, str] | None = None, +) -> bytes: """Generate a tarball containing SDK code and the user's prompt. The tarball contains: @@ -205,6 +233,10 @@ def _generate_tarball(prompt: str, repos: list[RepoSource] | None = None) -> byt _add_file_to_tar(tar, "prompt.txt", prompt) _add_file_to_tar(tar, "setup.sh", preset_files["setup.sh"], mode=0o755) + if callback_files: + for path, content in callback_files.items(): + _add_file_to_tar(tar, path, content) + # Add repos config if repos specified (SDK workspace handles cloning) if repos: repos_config = [r.model_dump(exclude_none=True) for r in repos] @@ -395,7 +427,10 @@ async def create_automation_from_prompt( model = resolve_model_profile_for_user(body.model, user) # 1. Generate tarball with SDK code, prompt, and optional repos config - tarball_content = _generate_tarball(body.prompt, repos=body.repos) + callbacks, callback_files = _materialize_callback_files(body.callbacks) + tarball_content = _generate_tarball( + body.prompt, repos=body.repos, callback_files=callback_files + ) # 2. Upload tarball to storage upload_id = uuid.uuid4() @@ -451,7 +486,10 @@ async def create_automation_from_prompt( setup_script_path="setup.sh", entrypoint=_get_preset_entrypoint(), timeout=body.timeout, - keep_alive=body.keep_alive, + keep_alive=body.keep_alive + if body.keep_alive is not None + else (True if callbacks else None), + callbacks=callbacks, ) session.add(automation) await session.flush() @@ -570,6 +608,10 @@ class CreatePluginAutomationRequest(BaseModel): "completion (or after post-run callbacks, when configured)." ), ) + callbacks: list[AutomationCallback] | None = Field( + default=None, + description="Post-run callbacks to run after COMPLETED or FAILED runs.", + ) repos: list[RepoSource] | None = Field( default=None, description=( @@ -658,6 +700,7 @@ def _generate_plugin_tarball( *, experiment_id: str | None = None, variants: list[ExperimentVariant] | None = None, + callback_files: dict[str, str] | None = None, ) -> bytes: """Generate a tarball containing SDK code, plugin config, and prompt. @@ -698,6 +741,10 @@ def _generate_plugin_tarball( tar, "plugins_config.json", json.dumps(plugins_config, indent=2) ) + if callback_files: + for path, content in callback_files.items(): + _add_file_to_tar(tar, path, content) + if repos: repos_config = [r.model_dump(exclude_none=True) for r in repos] _add_file_to_tar( @@ -747,12 +794,14 @@ async def create_automation_from_plugin( ) # 1. Generate tarball with SDK code, plugin/experiment config, and prompt + callbacks, callback_files = _materialize_callback_files(body.callbacks) tarball_content = _generate_plugin_tarball( body.plugins, body.prompt, repos=body.repos, experiment_id=body.experiment_id, variants=variants, + callback_files=callback_files, ) # 2. Upload tarball to storage @@ -819,7 +868,10 @@ async def create_automation_from_plugin( setup_script_path="setup.sh", entrypoint=_get_preset_entrypoint(), timeout=body.timeout, - keep_alive=body.keep_alive, + keep_alive=body.keep_alive + if body.keep_alive is not None + else (True if callbacks else None), + callbacks=callbacks, ) session.add(automation) await session.flush() diff --git a/openhands/automation/router.py b/openhands/automation/router.py index f5398d2f..6cd6dd0d 100644 --- a/openhands/automation/router.py +++ b/openhands/automation/router.py @@ -13,10 +13,16 @@ from sqlalchemy.orm import selectinload from openhands.automation.auth import AuthenticatedUser, require_permission +from openhands.automation.callbacks import ( + cleanup_run_after_callbacks_if_ready, + complete_callback_records, + schedule_and_dispatch_callbacks_for_run, +) from openhands.automation.db import get_session from openhands.automation.models import ( Automation, AutomationRun, + AutomationRunCallback, AutomationRunStatus, TarballUpload, ) @@ -24,8 +30,10 @@ from openhands.automation.schemas import ( AutomationListResponse, AutomationResponse, + AutomationRunCallbackResponse, AutomationRunListResponse, AutomationRunResponse, + CallbackCompleteRequest, CreateAutomationRequest, RunCompleteRequest, UpdateAutomationRequest, @@ -87,7 +95,12 @@ async def create_automation( setup_script_path=body.setup_script_path, entrypoint=body.entrypoint, timeout=body.timeout, - keep_alive=body.keep_alive, + keep_alive=body.keep_alive + if body.keep_alive is not None + else (True if body.callbacks else None), + callbacks=[c.model_dump(exclude_none=True) for c in body.callbacks] + if body.callbacks + else None, ) session.add(auto) await session.flush() @@ -153,6 +166,16 @@ async def update_automation( if "model" in update_data: update_data["model"] = resolve_model_profile_for_user(body.model, user) + if "callbacks" in update_data and body.callbacks is not None: + update_data["callbacks"] = [ + c.model_dump(exclude_none=True) for c in body.callbacks + ] + if ( + body.callbacks + and "keep_alive" not in update_data + and auto.keep_alive is None + ): + update_data["keep_alive"] = True original_prompt = auto.prompt for field, value in update_data.items(): @@ -393,9 +416,11 @@ async def complete_run( await session.refresh(run) logger.info("Run %s → %s", run_id, new_status.value) - # Clean up immediately when this automation owns explicit cleanup. Once - # post-run callbacks exist, this path should run them before deleting. - if run.sandbox_id and automation.keep_alive is not True: + callbacks_created = await schedule_and_dispatch_callbacks_for_run(session, run) + + # Clean up immediately when this automation owns explicit cleanup and no + # post-run callbacks need to run first. + if callbacks_created == 0 and run.sandbox_id and automation.keep_alive is not True: # Fire-and-forget sandbox deletion in background from openhands.automation.config import get_settings @@ -427,6 +452,64 @@ async def complete_run( return AutomationRunResponse.model_validate(run) +@router.post("/runs/{run_id}/callbacks/complete") +async def complete_run_callbacks( + run_id: uuid.UUID, + body: CallbackCompleteRequest, + user: AuthenticatedUser = Depends(_require_manage_automations), + session: AsyncSession = Depends(get_session), +) -> dict[str, str]: + """Record completion of a post-run callback chain.""" + result = await session.execute( + select(AutomationRun) + .where(AutomationRun.id == run_id) + .options(selectinload(AutomationRun.automation)) + ) + run = result.scalars().first() + if run is None: + raise HTTPException(status.HTTP_404_NOT_FOUND, detail="Run not found") + + automation = run.automation + if automation.user_id != user.user_id or automation.org_id != user.org_id: + raise HTTPException(status.HTTP_403_FORBIDDEN, detail="Not your automation") + + await complete_callback_records(session, run, body.callbacks) + await cleanup_run_after_callbacks_if_ready(session, run) + logger.info("Run %s callback chain completion recorded", run_id) + return {"status": "ok"} + + +@router.get("/runs/{run_id}/callbacks") +async def list_run_callbacks( + run_id: uuid.UUID, + user: AuthenticatedUser = Depends(_require_manage_automations), + session: AsyncSession = Depends(get_session), +) -> list[AutomationRunCallbackResponse]: + """List post-run callback executions for a run.""" + result = await session.execute( + select(AutomationRun) + .where(AutomationRun.id == run_id) + .options(selectinload(AutomationRun.automation)) + ) + run = result.scalars().first() + if run is None: + raise HTTPException(status.HTTP_404_NOT_FOUND, detail="Run not found") + + automation = run.automation + if automation.user_id != user.user_id or automation.org_id != user.org_id: + raise HTTPException(status.HTTP_403_FORBIDDEN, detail="Not your automation") + + callback_result = await session.execute( + select(AutomationRunCallback) + .where(AutomationRunCallback.run_id == run_id) + .order_by(AutomationRunCallback.order) + ) + return [ + AutomationRunCallbackResponse.model_validate(callback) + for callback in callback_result.scalars().all() + ] + + # --- Run cancellation --- diff --git a/openhands/automation/schemas.py b/openhands/automation/schemas.py index 5dbd1245..f1ade27f 100644 --- a/openhands/automation/schemas.py +++ b/openhands/automation/schemas.py @@ -6,7 +6,15 @@ from typing import Annotated, Literal from croniter import croniter -from pydantic import BaseModel, ConfigDict, Discriminator, Field, Tag, field_validator +from pydantic import ( + BaseModel, + ConfigDict, + Discriminator, + Field, + Tag, + field_validator, + model_validator, +) from openhands.automation.config import get_config from openhands.automation.constants import MODEL_PROFILE_PATTERN @@ -248,6 +256,72 @@ def _validate_command_string( return v +_CALLBACK_NAME_RE = re.compile(r"^[A-Za-z0-9_-]{1,100}$") + + +class AutomationCallback(BaseModel): + """Post-run callback configuration for an automation.""" + + model_config = ConfigDict(extra="forbid") + + name: str = Field(..., min_length=1, max_length=100) + on: list[Literal["COMPLETED", "FAILED"]] = Field( + default_factory=lambda: ["COMPLETED", "FAILED"], + min_length=1, + max_length=2, + description="Main run terminal statuses that trigger this callback.", + ) + entrypoint: str | None = Field( + default=None, + description="Command to run inside the extracted automation work directory.", + ) + inline_python: str | None = Field( + default=None, + min_length=1, + max_length=50000, + description="Inline Python source to execute as the callback.", + ) + timeout: int | None = Field( + default=None, + description="Maximum callback execution time in seconds.", + ) + + @field_validator("name") + @classmethod + def validate_name(cls, v: str) -> str: + if not _CALLBACK_NAME_RE.match(v): + raise ValueError( + "callback name must contain only letters, numbers, underscores, " + "or hyphens" + ) + return v + + @field_validator("on") + @classmethod + def validate_on(cls, v: list[str]) -> list[str]: + if len(set(v)) != len(v): + raise ValueError("callback statuses must be unique") + return v + + @field_validator("entrypoint") + @classmethod + def validate_entrypoint(cls, v: str | None) -> str | None: + return _validate_command_string(v, "entrypoint") + + @field_validator("timeout") + @classmethod + def validate_timeout(cls, v: int | None) -> int | None: + return _validate_timeout(v) + + @model_validator(mode="after") + def validate_exactly_one_source(self) -> "AutomationCallback": + if (self.entrypoint is None) == (self.inline_python is None): + raise ValueError( + "callback must provide exactly one of entrypoint or inline_python" + ) + return self + + # --- Requests --- @@ -291,6 +365,10 @@ class CreateAutomationRequest(BaseModel): "completion (or after post-run callbacks, when configured)." ), ) + callbacks: list[AutomationCallback] | None = Field( + default=None, + description="Post-run callbacks to run after COMPLETED or FAILED runs.", + ) @field_validator("tarball_path") @classmethod @@ -345,6 +423,7 @@ class UpdateAutomationRequest(BaseModel): entrypoint: str | None = Field(default=None) timeout: int | None = Field(default=None) keep_alive: bool | None = Field(default=None) + callbacks: list[AutomationCallback] | None = Field(default=None) enabled: bool | None = None @field_validator("tarball_path") @@ -598,6 +677,7 @@ class AutomationResponse(BaseModel): entrypoint: str timeout: int | None keep_alive: bool | None + callbacks: list[AutomationCallback] | None enabled: bool last_triggered_at: UtcDatetime | None created_at: UtcDatetime @@ -625,6 +705,51 @@ class RunCompleteRequest(BaseModel): error: str | None = None +class CallbackCompleteItem(BaseModel): + """Completion result for one post-run callback.""" + + model_config = ConfigDict(extra="forbid") + + id: uuid.UUID + name: str + status: Literal["COMPLETED", "FAILED"] + exit_code: int | None = None + stdout: str | None = Field(default=None, max_length=5000) + stderr: str | None = Field(default=None, max_length=5000) + error_detail: str | None = Field(default=None, max_length=5000) + started_at: UtcDatetime | None = None + completed_at: UtcDatetime | None = None + + +class CallbackCompleteRequest(BaseModel): + """Completion payload from the in-sandbox callback chain wrapper.""" + + model_config = ConfigDict(extra="forbid") + + callbacks: list[CallbackCompleteItem] = Field(default_factory=list) + + +class AutomationRunCallbackResponse(BaseModel): + """Response for a post-run callback execution.""" + + id: uuid.UUID + run_id: uuid.UUID + name: str + trigger_status: RunStatus + entrypoint: str + timeout: int | None + status: RunStatus + bash_command_id: str | None + error_detail: str | None + order: int + started_at: UtcDatetime | None + timeout_at: UtcDatetime | None + completed_at: UtcDatetime | None + created_at: UtcDatetime + + model_config = {"from_attributes": True} + + class AutomationRunResponse(BaseModel): """Response for a single automation run.""" diff --git a/openhands/automation/watchdog.py b/openhands/automation/watchdog.py index 2535490c..e08f272e 100644 --- a/openhands/automation/watchdog.py +++ b/openhands/automation/watchdog.py @@ -20,6 +20,7 @@ from sqlalchemy.orm import selectinload from openhands.automation.backends import get_backend +from openhands.automation.callbacks import schedule_and_dispatch_callbacks_for_run from openhands.automation.config import Settings from openhands.automation.models import ( Automation, @@ -91,6 +92,22 @@ async def _verify_and_mark_run( ) ) result: CursorResult = await session.execute(stmt) # type: ignore[assignment] + if result.rowcount > 0: + await session.refresh(run) + callbacks_created = await schedule_and_dispatch_callbacks_for_run( + session, run + ) + if callbacks_created == 0 and _should_cleanup_sandbox_after_terminal( + run, keep_alive + ): + try: + await backend.cleanup_after_verification(run_id) + except Exception as cleanup_error: + logger.warning( + "Cleanup after verification exception failed: %s", + cleanup_error, + extra=extra, + ) return result.rowcount > 0 if verification.verified: @@ -168,15 +185,22 @@ async def _verify_and_mark_run( ) result = await session.execute(stmt) # type: ignore[assignment] - if result.rowcount > 0 and _should_cleanup_sandbox_after_terminal( - run, keep_alive - ): - try: - await backend.cleanup_after_verification(run_id) - except Exception as e: - logger.warning( - "Cleanup after terminal verification failed: %s", e, extra=extra - ) + if result.rowcount > 0: + await session.refresh(run) + callbacks_created = await schedule_and_dispatch_callbacks_for_run( + session, run + ) + if callbacks_created == 0 and _should_cleanup_sandbox_after_terminal( + run, keep_alive + ): + try: + await backend.cleanup_after_verification(run_id) + except Exception as e: + logger.warning( + "Cleanup after terminal verification failed: %s", + e, + extra=extra, + ) return result.rowcount > 0 # Verification failed - execution environment not available or command still running @@ -187,14 +211,6 @@ async def _verify_and_mark_run( extra=extra, ) - # Clean up resources via backend only when the automation owns explicit - # cleanup. Otherwise, leave cleanup to the runtime TTL reaper. - if _should_cleanup_sandbox_after_terminal(run, keep_alive): - try: - await backend.cleanup_after_verification(run_id) - except Exception as e: - logger.warning("Cleanup after verification failed: %s", e, extra=extra) - error_msg = verification.error or "no completion callback received" logger.warning( @@ -219,6 +235,16 @@ async def _verify_and_mark_run( ) ) result = await session.execute(stmt) # type: ignore[assignment] + if result.rowcount > 0: + await session.refresh(run) + callbacks_created = await schedule_and_dispatch_callbacks_for_run(session, run) + if callbacks_created == 0 and _should_cleanup_sandbox_after_terminal( + run, keep_alive + ): + try: + await backend.cleanup_after_verification(run_id) + except Exception as e: + logger.warning("Cleanup after verification failed: %s", e, extra=extra) return result.rowcount > 0 diff --git a/tests/test_callbacks.py b/tests/test_callbacks.py new file mode 100644 index 00000000..32d8f445 --- /dev/null +++ b/tests/test_callbacks.py @@ -0,0 +1,291 @@ +"""Tests for post-run callback scheduling and completion.""" + +import uuid +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from sqlalchemy import select +from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine +from sqlalchemy.pool import StaticPool + +from openhands.automation.backends.base import ExecutionContext +from openhands.automation.callbacks import ( + cleanup_run_after_callbacks_if_ready, + complete_callback_records, + schedule_and_dispatch_callbacks_for_run, +) +from openhands.automation.models import ( + Automation, + AutomationRun, + AutomationRunCallback, + AutomationRunCallbackStatus, + AutomationRunStatus, + Base, +) +from openhands.automation.schemas import CallbackCompleteItem +from openhands.automation.utils import utcnow + + +TEST_USER_ID = uuid.UUID("12345678-1234-5678-1234-567812345678") +TEST_ORG_ID = uuid.UUID("87654321-4321-8765-4321-876543218765") + + +@pytest.fixture +async def async_session_factory(): + engine = create_async_engine( + "sqlite+aiosqlite:///:memory:", + connect_args={"check_same_thread": False}, + poolclass=StaticPool, + ) + async with engine.begin() as conn: + await conn.run_sync(Base.metadata.create_all) + yield async_sessionmaker(engine, class_=AsyncSession, expire_on_commit=False) + await engine.dispose() + + +@pytest.fixture +async def async_session(async_session_factory): + async with async_session_factory() as session: + yield session + + +def _automation(callbacks, keep_alive=None) -> Automation: + return Automation( + user_id=TEST_USER_ID, + org_id=TEST_ORG_ID, + name="Callback Automation", + trigger={"type": "cron", "schedule": "0 9 * * *", "timezone": "UTC"}, + tarball_path="s3://bucket/code.tar.gz", + entrypoint="python main.py", + keep_alive=keep_alive, + callbacks=callbacks, + ) + + +def _run(automation: Automation, status: AutomationRunStatus) -> AutomationRun: + return AutomationRun( + automation_id=automation.id, + status=status, + sandbox_id="sandbox-123", + completed_at=utcnow(), + ) + + +def _backend() -> MagicMock: + backend = MagicMock() + backend.get_existing_execution_context = AsyncMock( + return_value=ExecutionContext(agent_url="http://agent", session_key="key") + ) + backend.get_work_dir.return_value = "/workspace/project" + backend.cleanup_after_verification = AsyncMock() + return backend + + +async def _callback_rows(session: AsyncSession, run_id: uuid.UUID): + result = await session.execute( + select(AutomationRunCallback) + .where(AutomationRunCallback.run_id == run_id) + .order_by(AutomationRunCallback.order) + ) + return list(result.scalars().all()) + + +async def test_schedule_and_dispatch_callbacks_matches_terminal_status(async_session): + automation = _automation( + [ + { + "name": "on-success", + "on": ["COMPLETED"], + "entrypoint": "python callbacks/success.py", + "timeout": 30, + }, + { + "name": "on-failure", + "on": ["FAILED"], + "entrypoint": "python callbacks/failure.py", + }, + ] + ) + async_session.add(automation) + await async_session.flush() + run = _run(automation, AutomationRunStatus.COMPLETED) + async_session.add(run) + await async_session.flush() + + backend = _backend() + with ( + patch("openhands.automation.callbacks.get_backend", return_value=backend), + patch( + "openhands.automation.callbacks._start_bash", new_callable=AsyncMock + ) as mock_start, + ): + mock_start.return_value = "cmd-123" + created = await schedule_and_dispatch_callbacks_for_run(async_session, run) + + assert created == 1 + mock_start.assert_awaited_once() + callbacks = await _callback_rows(async_session, run.id) + assert len(callbacks) == 1 + assert callbacks[0].name == "on-success" + assert callbacks[0].status == AutomationRunCallbackStatus.RUNNING + assert callbacks[0].entrypoint == "python callbacks/success.py" + assert callbacks[0].bash_command_id == "cmd-123" + + +async def test_schedule_inline_python_callback_is_in_wrapper_command(async_session): + automation = _automation( + [ + { + "name": "inline", + "on": ["FAILED"], + "inline_python": "print('failed')\n", + } + ] + ) + async_session.add(automation) + await async_session.flush() + run = _run(automation, AutomationRunStatus.FAILED) + async_session.add(run) + await async_session.flush() + + backend = _backend() + with ( + patch("openhands.automation.callbacks.get_backend", return_value=backend), + patch( + "openhands.automation.callbacks._start_bash", new_callable=AsyncMock + ) as mock_start, + ): + mock_start.return_value = "cmd-inline" + created = await schedule_and_dispatch_callbacks_for_run(async_session, run) + + assert created == 1 + assert mock_start.await_args is not None + command = mock_start.await_args.args[3] + assert "failed" in command + + +async def test_schedule_and_dispatch_marks_callbacks_failed_when_start_fails( + async_session, +): + automation = _automation( + [ + { + "name": "notify", + "on": ["COMPLETED"], + "entrypoint": "python callbacks/notify.py", + } + ] + ) + async_session.add(automation) + await async_session.flush() + run = _run(automation, AutomationRunStatus.COMPLETED) + async_session.add(run) + await async_session.flush() + + backend = _backend() + with ( + patch("openhands.automation.callbacks.get_backend", return_value=backend), + patch( + "openhands.automation.callbacks._start_bash", new_callable=AsyncMock + ) as mock_start, + ): + mock_start.side_effect = RuntimeError("agent unavailable") + created = await schedule_and_dispatch_callbacks_for_run(async_session, run) + + assert created == 1 + callbacks = await _callback_rows(async_session, run.id) + assert callbacks[0].status == AutomationRunCallbackStatus.FAILED + assert callbacks[0].error_detail == "agent unavailable" + + +async def test_complete_callback_records_updates_status(async_session): + automation = _automation(callbacks=[], keep_alive=True) + async_session.add(automation) + await async_session.flush() + run = _run(automation, AutomationRunStatus.COMPLETED) + async_session.add(run) + await async_session.flush() + callback = AutomationRunCallback( + run_id=run.id, + name="notify", + trigger_status=AutomationRunStatus.COMPLETED, + entrypoint="python callbacks/notify.py", + status=AutomationRunCallbackStatus.RUNNING, + bash_command_id="cmd-123", + order=0, + ) + async_session.add(callback) + await async_session.flush() + + await complete_callback_records( + async_session, + run, + [ + CallbackCompleteItem( + id=callback.id, + name="notify", + status="COMPLETED", + exit_code=0, + ) + ], + ) + await async_session.flush() + + refreshed = await async_session.get(AutomationRunCallback, callback.id) + await async_session.refresh(refreshed) + assert refreshed.status == AutomationRunCallbackStatus.COMPLETED + assert refreshed.completed_at is not None + assert refreshed.error_detail is None + + +async def test_cleanup_after_callbacks_waits_for_unfinished_callbacks(async_session): + automation = _automation(callbacks=[], keep_alive=False) + async_session.add(automation) + await async_session.flush() + run = _run(automation, AutomationRunStatus.COMPLETED) + async_session.add(run) + await async_session.flush() + callback = AutomationRunCallback( + run_id=run.id, + name="notify", + trigger_status=AutomationRunStatus.COMPLETED, + entrypoint="python callbacks/notify.py", + status=AutomationRunCallbackStatus.RUNNING, + order=0, + ) + async_session.add(callback) + await async_session.flush() + + backend = _backend() + with patch("openhands.automation.callbacks.get_backend", return_value=backend): + cleaned = await cleanup_run_after_callbacks_if_ready(async_session, run) + + assert cleaned is False + backend.cleanup_after_verification.assert_not_called() + + +async def test_cleanup_after_callbacks_runs_when_all_callbacks_terminal(async_session): + automation = _automation(callbacks=[], keep_alive=False) + async_session.add(automation) + await async_session.flush() + run = _run(automation, AutomationRunStatus.COMPLETED) + async_session.add(run) + await async_session.flush() + callback = AutomationRunCallback( + run_id=run.id, + name="notify", + trigger_status=AutomationRunStatus.COMPLETED, + entrypoint="python callbacks/notify.py", + status=AutomationRunCallbackStatus.COMPLETED, + completed_at=utcnow(), + order=0, + ) + async_session.add(callback) + await async_session.flush() + + backend = _backend() + with patch("openhands.automation.callbacks.get_backend", return_value=backend): + cleaned = await cleanup_run_after_callbacks_if_ready(async_session, run) + + assert cleaned is True + backend.cleanup_after_verification.assert_awaited_once() diff --git a/tests/test_preset_router.py b/tests/test_preset_router.py index cd49f855..3b1996dc 100644 --- a/tests/test_preset_router.py +++ b/tests/test_preset_router.py @@ -15,9 +15,11 @@ _generate_plugin_tarball, _generate_tarball, _get_preset_entrypoint, + _materialize_callback_files, _replace_prompt_in_tarball, _resolve_experiment_variant_models, ) +from openhands.automation.schemas import AutomationCallback from openhands.sdk.plugin import PluginSource from openhands.workspace import RepoSource @@ -155,6 +157,43 @@ def test_generate_tarball_structure(self): # Note: load_skills.py and clone_repos.py are no longer needed # as the SDK workspace now provides these methods directly + def test_generate_tarball_with_callback_files(self): + """Generated prompt tarball can include materialized callback files.""" + tarball_bytes = _generate_tarball( + "Test prompt", + callback_files={"callbacks/notify.py": "print('done')\n"}, + ) + + with tarfile.open(fileobj=io.BytesIO(tarball_bytes), mode="r:gz") as tar: + names = tar.getnames() + assert "callbacks/notify.py" in names + callback_file = tar.extractfile("callbacks/notify.py") + assert callback_file is not None + assert callback_file.read().decode("utf-8") == "print('done')\n" + + def test_materialize_inline_callback_files(self): + """Inline preset callbacks are written to files and normalized.""" + callbacks, files = _materialize_callback_files( + [ + AutomationCallback( + name="mark_review_failed", + on=["FAILED"], + inline_python="print('failed')\n", + timeout=60, + ) + ] + ) + + assert files == {"callbacks/mark_review_failed.py": "print('failed')\n"} + assert callbacks == [ + { + "name": "mark_review_failed", + "on": ["FAILED"], + "timeout": 60, + "entrypoint": "python callbacks/mark_review_failed.py", + } + ] + def test_generate_tarball_prompt_content(self): """Generated tarball contains the user's prompt.""" prompt = "Write a Python script that prints 'Hello, World!'" diff --git a/tests/test_schemas.py b/tests/test_schemas.py index 28c9eb8c..c6d784f9 100644 --- a/tests/test_schemas.py +++ b/tests/test_schemas.py @@ -11,7 +11,11 @@ from datetime import UTC, datetime, timedelta, timezone from typing import Any +import pytest +from pydantic import ValidationError + from openhands.automation.schemas import ( + AutomationCallback, AutomationResponse, AutomationRunResponse, RunStatus, @@ -111,6 +115,7 @@ def _make_automation(self, **overrides: Any) -> AutomationResponse: entrypoint="python main.py", timeout=None, keep_alive=True, + callbacks=None, enabled=True, last_triggered_at=_NAIVE, created_at=_NAIVE, @@ -135,3 +140,47 @@ def test_none_last_triggered_at_remains_none(self): automation = self._make_automation(last_triggered_at=None) data = automation.model_dump(mode="json") assert data["last_triggered_at"] is None + + +class TestAutomationCallbackValidation: + def test_entrypoint_callback_is_valid(self): + callback = AutomationCallback( + name="notify_failure", + on=["FAILED"], + entrypoint="python callbacks/on_failure.py", + timeout=60, + ) + + assert callback.entrypoint == "python callbacks/on_failure.py" + assert callback.inline_python is None + + def test_inline_python_callback_is_valid(self): + callback = AutomationCallback( + name="mark_review_failed", + on=["FAILED"], + inline_python="print('failed')\n", + ) + + assert callback.inline_python == "print('failed')\n" + assert callback.entrypoint is None + + def test_callback_rejects_both_code_sources(self): + with pytest.raises(ValidationError, match="exactly one"): + AutomationCallback( + name="bad", + on=["FAILED"], + entrypoint="python callbacks/bad.py", + inline_python="print('bad')\n", + ) + + def test_callback_rejects_neither_code_source(self): + with pytest.raises(ValidationError, match="exactly one"): + AutomationCallback(name="bad", on=["FAILED"]) + + def test_callback_rejects_duplicate_statuses(self): + with pytest.raises(ValidationError, match="unique"): + AutomationCallback( + name="bad", + on=["FAILED", "FAILED"], + entrypoint="python callbacks/bad.py", + )