tests_python/test_swarm_runtime.py

from arena.codebase_archive import LocalCodebaseArchiveStore
from arena.model_catalog import ModelProviderConfig
from arena.run_models import Run
from arena.swarm_runtime import SwarmRuntime

from conftest import make_local_session


def test_start_run_seeds_prompt_first_context() -> None:
    runtime = SwarmRuntime(
        session_factory=make_local_session,
        archive_store=LocalCodebaseArchiveStore(),
    )
    run = Run(id="run-1", prompt="Build a CLI", rejoin_token="t" * 24)

    state = runtime.start_run(run)

    read_result = state.session.backend.read("/context/run.md")
    content = read_result.file_data.get("content", "")
    assert "Build a CLI" in content
    assert "/workspace" in content
    assert "handoff_spec.md" in content
    assert "field_notes.md" in content


def test_execute_run_creates_archive_and_returns_summary() -> None:
    archive_store = LocalCodebaseArchiveStore()
    runtime = SwarmRuntime(
        session_factory=make_local_session,
        archive_store=archive_store,
    )
    run = Run(id="run-1", prompt="Build a CLI", rejoin_token="t" * 24)
    state = runtime.start_run(run)
    state.session.backend.write("/workspace/app.py", "print('hello')")

    result = runtime.execute_run("run-1")

    assert result.summary == "Built the app."
    assert result.tasks[0].status == "completed"
    assert result.codebase_archive.file_count >= 1
    assert result.codebase_archive.size_bytes > 0
    assert result.codebase_archive.pointer.startswith("local-snapshot://")


def test_execute_run_respects_provider_workspace_dir() -> None:
    archive_store = LocalCodebaseArchiveStore()
    runtime = SwarmRuntime(
        session_factory=make_local_session,
        archive_store=archive_store,
    )
    run = Run(id="run-1", prompt="Build a CLI", rejoin_token="t" * 24)
    state = runtime.start_run(run)
    state.session.backend.write("/workspace/answer.txt", "42")

    result = runtime.execute_run("run-1")

    archive = archive_store.load(result.codebase_archive.pointer)
    assert any(f.path == "answer.txt" for f in archive.files)


def test_execute_run_closes_session_after_archiving() -> None:
    closed: list[str] = []

    def session_factory(thread_id: str | None = None):
        session = make_local_session(thread_id)
        session.close = lambda: closed.append(session.thread_id)
        return session

    runtime = SwarmRuntime(
        session_factory=session_factory,
        archive_store=LocalCodebaseArchiveStore(),
        sandbox_ttl_seconds=60,
    )
    run = Run(id="run-1", prompt="Build a CLI", rejoin_token="t" * 24)
    state = runtime.start_run(run)
    state.session.backend.write("/workspace/app.py", "print('hello')")

    runtime.execute_run("run-1")

    assert closed == ["run-1"]
    assert "run-1" not in runtime._runs


def test_sandbox_ttl_closes_unfinished_session() -> None:
    import time

    closed: list[str] = []

    def session_factory(thread_id: str | None = None):
        session = make_local_session(thread_id)
        session.close = lambda: closed.append(session.thread_id)
        return session

    runtime = SwarmRuntime(
        session_factory=session_factory,
        archive_store=LocalCodebaseArchiveStore(),
        sandbox_ttl_seconds=0.01,
    )
    run = Run(id="run-1", prompt="Build a CLI", rejoin_token="t" * 24)

    runtime.start_run(run)
    deadline = time.time() + 1
    while time.time() < deadline and not closed:
        time.sleep(0.01)

    assert closed == ["run-1"]
    assert "run-1" not in runtime._runs


def test_touch_run_resets_sandbox_ttl() -> None:
    import time

    closed: list[str] = []

    def session_factory(thread_id: str | None = None):
        session = make_local_session(thread_id)
        session.close = lambda: closed.append(session.thread_id)
        return session

    runtime = SwarmRuntime(
        session_factory=session_factory,
        archive_store=LocalCodebaseArchiveStore(),
        sandbox_ttl_seconds=0.06,
    )
    run = Run(id="run-1", prompt="Build a CLI", rejoin_token="t" * 24)

    runtime.start_run(run)
    time.sleep(0.03)
    assert runtime.touch_run("run-1") is True
    time.sleep(0.04)
    assert closed == []

    deadline = time.time() + 1
    while time.time() < deadline and not closed:
        time.sleep(0.01)

    assert closed == ["run-1"]


def test_start_run_passes_model_provider_to_session_factory() -> None:
    seen: list[ModelProviderConfig | None] = []

    def session_factory(thread_id: str | None = None, model_provider: ModelProviderConfig | None = None):
        seen.append(model_provider)
        return make_local_session(thread_id)

    runtime = SwarmRuntime(
        session_factory=session_factory,
        archive_store=LocalCodebaseArchiveStore(),
    )
    config = ModelProviderConfig(provider="custom", model="local-model", api_key="secret")
    run = Run(id="run-1", prompt="Build a CLI", rejoin_token="t" * 24)

    runtime.start_run(run, model_provider=config)

    assert seen == [config]
    runtime.close_run("run-1")