From 25958736796ddb581767b511b77ed93496399aa2 Mon Sep 17 00:00:00 2001 From: Emile Riberdy Date: Mon, 7 Sep 2026 19:56:25 +0000 Subject: [PATCH 1/5] test(predict-rlm): consolidate core regression suite --- .github/workflows/tests.yml | 5 +- CHANGELOG.md | 7 + Makefile | 4 +- README.md | 33 + pyproject.toml | 7 +- tests/codex_lm/conftest.py | 29 +- tests/codex_lm/test_aforward.py | 45 - tests/codex_lm/test_aliases.py | 16 - tests/codex_lm/test_assemble.py | 55 - tests/codex_lm/test_auth.py | 323 +- tests/codex_lm/test_auto_retry.py | 198 +- tests/codex_lm/test_build_request.py | 508 +-- tests/codex_lm/test_cache_hit_history.py | 102 - tests/codex_lm/test_cli.py | 960 +---- tests/codex_lm/test_concurrent_cache.py | 86 +- tests/codex_lm/test_cost.py | 78 - tests/codex_lm/test_events.py | 46 - tests/codex_lm/test_forward.py | 148 +- tests/codex_lm/test_proxy_and_service_tier.py | 104 - tests/codex_lm/test_reasoning_effort.py | 151 - tests/codex_lm/test_retry_jitter.py | 118 - tests/codex_lm/test_stream_errors.py | 93 +- tests/codex_lm/test_stream_heartbeat.py | 133 +- tests/codex_lm/test_stream_redaction.py | 35 + tests/codex_lm/test_stream_timing_debug.py | 285 -- tests/codex_lm/test_terminal_completion.py | 65 - tests/codex_lm/test_usage.py | 392 +- tests/codex_lm/test_usage_tracker_hook.py | 166 - tests/codex_lm/test_ws_lm.py | 226 +- .../debian13-no-python/Dockerfile | 21 - .../python311-slim/Dockerfile | 17 - .../tbench-python-313/Dockerfile | 16 - .../ubuntu24-python-pip-no-venv/Dockerfile | 21 - tests/runtime_contracts/backends.py | 321 +- tests/runtime_contracts/conftest.py | 27 +- .../runtime_contracts/test_backend_matrix.py | 34 - .../runtime_contracts/test_error_contract.py | 25 - .../test_execution_contract.py | 127 +- tests/runtime_contracts/test_file_contract.py | 31 - .../runtime_contracts/test_submit_contract.py | 31 - .../test_timeout_contract.py | 26 - tests/runtime_contracts/test_tool_contract.py | 282 +- tests/test_adapter_contracts.py | 323 +- tests/test_bootstrap_controller.py | 24 - tests/test_callbacks.py | 202 +- tests/test_codex_ws_lm.py | 125 +- tests/test_debug.py | 84 - tests/test_direct_python_backend.py | 22 +- tests/test_empty_code_retry.py | 276 +- tests/test_exception_race.py | 518 --- .../test_external_input_adapter_contracts.py | 54 +- tests/test_file_sync.py | 662 +--- tests/test_files.py | 1077 +----- tests/test_in_context.py | 501 +-- tests/test_interpreter.py | 3257 ++--------------- tests/test_interpreter_io.py | 975 +---- tests/test_interpreter_unit.py | 345 -- tests/test_iteration_execution_timeout.py | 362 +- tests/test_iteration_usage.py | 62 - tests/test_jspi_async_operations.py | 226 +- tests/test_jspi_workspace_lifecycle.py | 77 - tests/test_lm_config.py | 39 - tests/test_logging.py | 158 - tests/test_predict_rlm.py | 3035 ++------------- tests/test_pydantic_fixes.py | 973 ----- tests/test_read_line_raw_hang.py | 236 -- tests/test_response_id_resync.py | 194 +- tests/test_rlm_gepa.py | 3086 ++-------------- tests/test_rlm_gepa_patch_merge_costs.py | 60 +- tests/test_rlm_skill_docs.py | 28 - tests/test_rlm_skills.py | 170 +- tests/test_runtime_hooks.py | 38 +- tests/test_sandbox_backend_benchmark.py | 209 -- tests/test_sbx_interpreter.py | 3245 +--------------- tests/test_sbx_pool.py | 358 +- tests/test_shared.py | 209 -- tests/test_skill_package_integration.py | 46 +- tests/test_small_kernel.py | 2029 +--------- tests/test_supervisor_client.py | 56 +- tests/test_telemetry.py | 163 +- tests/test_telemetry_analyzer.py | 122 +- tests/test_terminal_bench_web_search.py | 109 - tests/test_tool_call_timeout.py | 73 +- tests/test_trace.py | 837 +---- tests/test_trace_on_cancellation.py | 272 +- tests/test_triple_quote_preprocess.py | 87 - tests/test_workspace.py | 371 +- tests/test_wt_hooks.py | 28 - 88 files changed, 2750 insertions(+), 28050 deletions(-) delete mode 100644 tests/codex_lm/test_aforward.py delete mode 100644 tests/codex_lm/test_aliases.py delete mode 100644 tests/codex_lm/test_assemble.py delete mode 100644 tests/codex_lm/test_cache_hit_history.py delete mode 100644 tests/codex_lm/test_cost.py delete mode 100644 tests/codex_lm/test_events.py delete mode 100644 tests/codex_lm/test_proxy_and_service_tier.py delete mode 100644 tests/codex_lm/test_reasoning_effort.py delete mode 100644 tests/codex_lm/test_retry_jitter.py create mode 100644 tests/codex_lm/test_stream_redaction.py delete mode 100644 tests/codex_lm/test_stream_timing_debug.py delete mode 100644 tests/codex_lm/test_terminal_completion.py delete mode 100644 tests/codex_lm/test_usage_tracker_hook.py delete mode 100644 tests/fixtures/bootstrap_controller/debian13-no-python/Dockerfile delete mode 100644 tests/fixtures/bootstrap_controller/python311-slim/Dockerfile delete mode 100644 tests/fixtures/bootstrap_controller/tbench-python-313/Dockerfile delete mode 100644 tests/fixtures/bootstrap_controller/ubuntu24-python-pip-no-venv/Dockerfile delete mode 100644 tests/runtime_contracts/test_backend_matrix.py delete mode 100644 tests/runtime_contracts/test_error_contract.py delete mode 100644 tests/runtime_contracts/test_file_contract.py delete mode 100644 tests/runtime_contracts/test_submit_contract.py delete mode 100644 tests/runtime_contracts/test_timeout_contract.py delete mode 100644 tests/test_debug.py delete mode 100644 tests/test_exception_race.py delete mode 100644 tests/test_interpreter_unit.py delete mode 100644 tests/test_iteration_usage.py delete mode 100644 tests/test_jspi_workspace_lifecycle.py delete mode 100644 tests/test_logging.py delete mode 100644 tests/test_pydantic_fixes.py delete mode 100644 tests/test_read_line_raw_hang.py delete mode 100644 tests/test_rlm_skill_docs.py delete mode 100644 tests/test_sandbox_backend_benchmark.py delete mode 100644 tests/test_shared.py delete mode 100644 tests/test_terminal_bench_web_search.py delete mode 100644 tests/test_triple_quote_preprocess.py delete mode 100644 tests/test_wt_hooks.py diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index af4b2c20..030799eb 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -7,9 +7,8 @@ on: branches: [main] jobs: - # Core package, no optional extras installed. Proves predict_rlm and its core - # logic import and pass standalone. Optional-subsystem modules (sbx/gepa/codex_lm) - # are deselected here, so a green run is a real "works without extras" signal. + # Core package plus local CPython subprocess contracts, without optional + # extras, Deno execution, or external sandbox services. unit-core: runs-on: ubuntu-latest strategy: diff --git a/CHANGELOG.md b/CHANGELOG.md index 0692f6b7..6aeb560a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,13 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Changed + +- Reduced the package regression suite to behavioral contracts; consolidated + shared backend execution tests and removed duplicate inherited cases, + static/presentation checks, and one-off benchmark/example checks. +- Direct CPython contracts now run in the core tier without the SBX extra. + ### Breaking Changes - `SbxBackend.shutdown()` can no longer be called from the event loop that owns diff --git a/Makefile b/Makefile index 8f755375..ea1b162f 100644 --- a/Makefile +++ b/Makefile @@ -22,8 +22,8 @@ test-integration: # --- Per-suite targets (one per CI job; each installs only what it needs) --- -# Core: no extras. Pure host-side logic; proves the package works standalone -# (passes with neither Deno nor websockets present). +# Core: no extras. Host logic and local CPython subprocess contracts; +# no Deno execution, websockets, or external sandbox service. test-core: uv run pytest -m "not integration and not sbx and not gepa and not codex_lm" $(ARGS) diff --git a/README.md b/README.md index 571ee3c2..b68607bc 100644 --- a/README.md +++ b/README.md @@ -354,6 +354,39 @@ result = await rlm.acall(query="...") Register globally instead with `dspy.configure(callbacks=[...])` and the same handlers fire for every `PredictRLM` instance. +## Tests + +The default package regression suite is `tests/`; example-local suites are +separate. Tests protect observable behavior, not exports, defaults, prompt wording, +docstrings, source text, or mocked argument forwarding. + +- `tests/runtime_contracts/`: shared execution, state, submission, files, tools, + and recovery contracts across Direct, JSPI, and SBX. Backend-specific gaps are + explicit in `backends.py`; native Direct callbacks are serial, and deferred + submission is a Direct-only interpreter API. +- RLM, adapter, file, and workspace tests: generated-code execution, typed + predictions, input/output handling, cancellation ownership, and data-loss boundaries. +- Trace and telemetry tests: evidence projection, accounting, redaction, and + failure classification. +- GEPA tests: candidate acceptance, merge guardrails, evaluation artifacts, + resume, and spend. +- `tests/codex_lm/`: auth, transport completion/recovery, caching, and usage. + +```bash +make test-unit # all extras; no Deno or external SBX execution +make test-integration-jspi # real Deno/Pyodide contracts; no LM credentials +make test-core # no extras; includes local CPython processes +make test-sbx # local WebSocket supervisor and pool contracts +make test-gepa +make test-codex-lm +``` + +Real SBX tests require the CLI, login, and `make test-integration-sbx`. +Bootstrap image tests additionally require Docker and +`PREDICT_RLM_RUN_BOOTSTRAP_DOCKER_TESTS=1`. Timing-sensitive `local` tests run +locally but are excluded from CI. Add a case only for a distinct failure or +invariant; extend the shared contract instead of copying it into backend suites. + ## Next steps - [Custom path inputs](docs/custom-path-inputs.md) — add file-, workspace-, or diff --git a/pyproject.toml b/pyproject.toml index f8622a8b..79638de0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -72,10 +72,9 @@ addopts = "-s -rP --testdox" asyncio_mode = "auto" testpaths = ["tests"] markers = [ - # Tier axis: a test that needs a real backend process (real Deno/WASM, or real - # Docker Sandboxes). Combine with `sbx` for real-SBX tests; `integration and not - # sbx` is the real-Deno/JSPI tier. - "integration: requires a real backend process (real Deno/WASM sandbox, or real SBX)", + # External runtime axis: Deno/WASM or real Docker Sandboxes. Local CPython + # subprocess contracts run in core; local WebSocket supervisors use `sbx`. + "integration: requires Deno/WASM or the real Docker Sandboxes service", # Install axis: needs the [sbx] extra (websockets) and exercises the supervisor/SBX # backend. `sbx and not integration` runs against a local runner/supervisor seam # (no real SBX); `sbx and integration` needs the real Docker Sandboxes service. diff --git a/tests/codex_lm/conftest.py b/tests/codex_lm/conftest.py index 7a686af7..0356c9b1 100644 --- a/tests/codex_lm/conftest.py +++ b/tests/codex_lm/conftest.py @@ -1,4 +1,3 @@ -import json from pathlib import Path from types import SimpleNamespace from typing import Any @@ -32,33 +31,17 @@ def reset_dspy_cache(): @pytest.fixture(autouse=True) def _disable_codex_retries_in_tests(monkeypatch): - """By default, tests see CodexLM behave like the pre-tenacity version: - one attempt, no retry, no backoff. Stream-error tests rely on a single - ``iter(events)`` side-effect that would be exhausted on retry. Tests - that want to exercise retry behaviour (``tests/test_auto_retry.py``) - re-raise this fixture's values to enable retries with zero backoff. - """ + """Retries are opt-in; tests never sleep through production backoff.""" monkeypatch.setattr("dspy_codex_lm.lm.CODEX_STREAM_MAX_ATTEMPTS", 1) monkeypatch.setattr("dspy_codex_lm.lm.CODEX_STREAM_WAIT_MULTIPLIER", 0.0) monkeypatch.setattr("dspy_codex_lm.lm.CODEX_STREAM_WAIT_MAX", 0.0) -@pytest.fixture -def fake_auth_file(tmp_path: Path) -> Path: - path = tmp_path / "auth.json" - path.write_text( - json.dumps( - { - "tokens": { - "access_token": "fake-access-token", - "account_id": "fake-account-id", - "refresh_token": "fake-refresh", - "id_token": "fake-id", - } - } - ) - ) - return path +@pytest.fixture(autouse=True) +def isolate_auth_home(tmp_path, monkeypatch): + monkeypatch.setattr(Path, "home", classmethod(lambda cls: tmp_path)) + monkeypatch.delenv("CODEX_LM_AUTH_PROFILE", raising=False) + monkeypatch.delenv("CODEX_LM_ENABLE_LEGACY_AUTH_FALLBACK", raising=False) @pytest.fixture diff --git a/tests/codex_lm/test_aforward.py b/tests/codex_lm/test_aforward.py deleted file mode 100644 index 73ade857..00000000 --- a/tests/codex_lm/test_aforward.py +++ /dev/null @@ -1,45 +0,0 @@ -from unittest import mock - -import dspy -from conftest import build_stream_events - - -async def _async_iter(items): - for item in items: - yield item - - -def _patch_aresponses(events, call_count=None): - async def fake(**_): - if call_count is not None: - call_count["n"] += 1 - return _async_iter(events) - - return mock.patch("dspy_codex_lm.lm.litellm.aresponses", side_effect=fake) - - -async def test_aforward_single_call(lm): - events = build_stream_events("4", input_tokens=50, output_tokens=1) - with _patch_aresponses(events): - resp = await lm.aforward(prompt="What is 2+2?") - assert resp.output[0].content[0].text == "4" - assert resp.usage.cost > 0 - - -async def test_aforward_via_dspy_predict(lm): - text = "[[ ## answer ## ]]\n4\n[[ ## completed ## ]]" - events = build_stream_events(text, input_tokens=50, output_tokens=10) - with _patch_aresponses(events): - dspy.configure(lm=lm) - result = await dspy.Predict("question -> answer").acall(question="What is 2+2?") - assert result.answer.strip() == "4" - - -async def test_aforward_second_call_hits_cache(lm): - events = build_stream_events("first", input_tokens=10, output_tokens=3) - count = {"n": 0} - with _patch_aresponses(events, call_count=count): - r1 = await lm.aforward(prompt="same") - r2 = await lm.aforward(prompt="same") - assert count["n"] == 1 - assert r1.output[0].content[0].text == r2.output[0].content[0].text == "first" diff --git a/tests/codex_lm/test_aliases.py b/tests/codex_lm/test_aliases.py deleted file mode 100644 index a7a2ab42..00000000 --- a/tests/codex_lm/test_aliases.py +++ /dev/null @@ -1,16 +0,0 @@ -from dspy_codex_lm import CodexHTTPLM, CodexLM, CodexWSLM - - -def test_codex_lm_aliases_websocket_by_default(): - assert CodexLM is CodexWSLM - - -def test_http_transport_remains_available_as_codex_http_lm(): - lm = CodexHTTPLM( - model="gpt-5.3-codex", - access_token="fake-access", - account_id="fake-account", - ) - - assert isinstance(lm, CodexHTTPLM) - assert not isinstance(lm, CodexWSLM) diff --git a/tests/codex_lm/test_assemble.py b/tests/codex_lm/test_assemble.py deleted file mode 100644 index 84c518cb..00000000 --- a/tests/codex_lm/test_assemble.py +++ /dev/null @@ -1,55 +0,0 @@ -from conftest import make_completed, make_text_delta -from litellm.types.llms.openai import ResponsesAPIResponse - - -def _run_events(lm, events): - state = lm._fresh_state() - for ev in events: - lm._handle_event(ev, state) - return lm._assemble( - text_parts=state["text_parts"], - usage_raw=state["usage_raw"], - response_id=state["response_id"], - model_name=state["model_name"], - ) - - -def test_assemble_returns_responses_api_response(lm): - events = [make_text_delta("hi"), make_completed(input_tokens=5, output_tokens=2)] - response = _run_events(lm, events) - assert isinstance(response, ResponsesAPIResponse) - - -def test_assembled_text_is_concatenated(lm): - events = [ - make_text_delta("he"), - make_text_delta("llo"), - make_text_delta(" world"), - make_completed(), - ] - response = _run_events(lm, events) - assert response.output[0].content[0].text == "hello world" - - -def test_assembled_cost_matches_rate_card(lm): - events = [ - make_text_delta("x"), - make_completed(input_tokens=100, output_tokens=10, model="gpt-5.3-codex"), - ] - response = _run_events(lm, events) - expected = (100 * 1.75e-6) + (10 * 1.4e-5) - assert abs(response.usage.cost - expected) < 1e-12 - assert abs(response._hidden_params["response_cost"] - expected) < 1e-12 - - -def test_assemble_with_no_usage_defaults_to_zero_cost(lm): - state = lm._fresh_state() - # only text, never a completed event - lm._handle_event(make_text_delta("x"), state) - response = lm._assemble( - text_parts=state["text_parts"], - usage_raw=state["usage_raw"], - response_id=state["response_id"], - model_name=state["model_name"], - ) - assert response.usage.cost == 0.0 diff --git a/tests/codex_lm/test_auth.py b/tests/codex_lm/test_auth.py index 84b228a6..ed374da4 100644 --- a/tests/codex_lm/test_auth.py +++ b/tests/codex_lm/test_auth.py @@ -6,7 +6,6 @@ from dspy_codex_lm.auth import ( CODEX_LM_AUTH_PROFILE_ENV, auth_status_metadata, - clear_active_profile, enable_auth_profile, get_active_profile, import_auth_profile, @@ -14,7 +13,6 @@ list_auth_profiles, list_enabled_auth_profiles, load_codex_auth, - profile_auth_path, remove_auth_profile, set_active_profile, validate_profile_name, @@ -25,7 +23,7 @@ @pytest.fixture -def fake_home(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Path: +def fake_home(tmp_path: Path, monkeypatch: pytest.MonkeyPatch, isolate_auth_home) -> Path: monkeypatch.setattr(Path, "home", classmethod(lambda cls: tmp_path)) monkeypatch.delenv(CODEX_LM_AUTH_PROFILE_ENV, raising=False) monkeypatch.delenv(LEGACY_AUTH_FALLBACK_ENV, raising=False) @@ -56,25 +54,7 @@ def _write_auth( return path -def test_loads_access_and_account_id(tmp_path: Path): - auth = tmp_path / "auth.json" - _write_auth(auth, account_id="acct-1") - assert load_codex_auth(auth) == ("abc", "acct-1") - - -def test_missing_file_raises(tmp_path: Path): - with pytest.raises(FileNotFoundError): - load_codex_auth(tmp_path / "nope.json") - - -def test_missing_fields_raises(tmp_path: Path): - auth = tmp_path / "auth.json" - auth.write_text(json.dumps({"tokens": {"access_token": "abc"}})) - with pytest.raises(KeyError): - load_codex_auth(auth) - - -def test_load_without_selection_does_not_fall_back_by_default(fake_home: Path): +def test_legacy_credentials_require_explicit_opt_in(fake_home, monkeypatch): _write_auth( fake_home / ".codex" / "auth.json", access_token="fallback-token", @@ -85,69 +65,15 @@ def test_load_without_selection_does_not_fall_back_by_default(fake_home: Path): load_codex_auth() message = str(excinfo.value) - assert "codex-lm auth login NAME" in message - assert "codex-lm auth import NAME" in message - assert "codex-lm auth use NAME" in message - assert LEGACY_AUTH_FALLBACK_ENV in message assert "fallback-token" not in message assert "acct-fallback" not in message - -def test_load_falls_back_to_codex_cli_auth_when_enabled( - fake_home: Path, - monkeypatch: pytest.MonkeyPatch, -): - _write_auth( - fake_home / ".codex" / "auth.json", - access_token="fallback-token", - account_id="acct-fallback", - ) monkeypatch.setenv(LEGACY_AUTH_FALLBACK_ENV, "1") assert load_codex_auth() == ("fallback-token", "acct-fallback") -def test_import_list_use_remove_profile(fake_home: Path, tmp_path: Path): - source = _write_auth( - tmp_path / "source-auth.json", - access_token="profile-token", - account_id="acct-profile", - ) - - dest = import_auth_profile("work.1", source) - - assert dest == profile_auth_path("work.1") - assert list_auth_profiles() == ["work.1"] - assert get_active_profile() == "work.1" - assert stat.S_IMODE(dest.stat().st_mode) == 0o600 - - set_active_profile("work.1") - assert get_active_profile() == "work.1" - assert load_codex_auth() == ("profile-token", "acct-profile") - - remove_auth_profile("work.1") - assert list_auth_profiles() == [] - assert get_active_profile() is None - - -def test_import_profile_sets_default_only_when_none( - fake_home: Path, - tmp_path: Path, -): - import_auth_profile( - "first", - _write_auth(tmp_path / "first.json", access_token="first-token"), - ) - import_auth_profile( - "second", - _write_auth(tmp_path / "second.json", access_token="second-token"), - ) - - assert get_active_profile() == "first" - assert load_codex_auth() == ("first-token", "acct-1234567890") - - -def test_import_display_profile_uses_safe_slug_and_metadata( +def test_profile_import_persists_private_credentials_and_removal_clears_selection( fake_home: Path, tmp_path: Path, ): @@ -157,21 +83,13 @@ def test_import_display_profile_uses_safe_slug_and_metadata( account_id="acct-profile", ) display_name = "gabriel@trampoline.ai" - slug = "gabriel-trampoline.ai" dest = import_auth_profile(display_name, source) + assert stat.S_IMODE(dest.stat().st_mode) == 0o600 - assert dest == fake_home / ".codex-lm" / "auth" / slug / "auth.json" - assert profile_auth_path(display_name) == dest assert list_auth_profiles() == [display_name] - assert json.loads((dest.parent / "profile.json").read_text(encoding="utf-8")) == { - "name": display_name, - "slug": slug, - } assert get_active_profile() == display_name - set_active_profile(display_name) - assert get_active_profile() == display_name assert load_codex_auth() == ("profile-token", "acct-profile") remove_auth_profile(display_name) @@ -180,23 +98,6 @@ def test_import_display_profile_uses_safe_slug_and_metadata( assert get_active_profile() is None -def test_legacy_profile_dirs_without_metadata_still_work( - fake_home: Path, -): - auth = _write_auth( - fake_home / ".codex-lm" / "auth" / "legacy" / "auth.json", - access_token="legacy-token", - account_id="acct-legacy", - ) - - assert list_auth_profiles() == ["legacy"] - assert profile_auth_path("legacy") == auth - - set_active_profile("legacy") - assert get_active_profile() == "legacy" - assert load_codex_auth() == ("legacy-token", "acct-legacy") - - def test_disable_enable_profile_persists_state_without_deleting_auth( fake_home: Path, tmp_path: Path, @@ -227,17 +128,6 @@ def test_disable_enable_profile_persists_state_without_deleting_auth( assert metadata["disabled"] is False -def test_legacy_profile_dirs_are_enabled_by_default(fake_home: Path): - _write_auth( - fake_home / ".codex-lm" / "auth" / "legacy" / "auth.json", - access_token="legacy-token", - account_id="acct-legacy", - ) - - assert is_auth_profile_disabled("legacy") is False - assert list_enabled_auth_profiles() == ["legacy"] - - def test_import_profile_rejects_slug_collision(fake_home: Path, tmp_path: Path): import_auth_profile( "a@b", @@ -253,35 +143,6 @@ def test_import_profile_rejects_slug_collision(fake_home: Path, tmp_path: Path): assert load_codex_auth(profile="a@b") == ("first-token", "acct-1234567890") -def test_env_profile_overrides_active_without_changing_marker( - fake_home: Path, - tmp_path: Path, - monkeypatch: pytest.MonkeyPatch, -): - import_auth_profile( - "active", - _write_auth( - tmp_path / "active.json", - access_token="active-token", - account_id="acct-active", - ), - ) - import_auth_profile( - "env", - _write_auth( - tmp_path / "env.json", - access_token="env-token", - account_id="acct-env", - ), - ) - set_active_profile("active") - - monkeypatch.setenv(CODEX_LM_AUTH_PROFILE_ENV, "env") - - assert load_codex_auth() == ("env-token", "acct-env") - assert get_active_profile() == "active" - - def test_disabled_explicit_env_and_active_profiles_fail_clearly( fake_home: Path, tmp_path: Path, @@ -310,67 +171,6 @@ def test_disabled_explicit_env_and_active_profiles_fail_clearly( load_codex_auth() -def test_explicit_path_ignores_profile_disabled_state( - fake_home: Path, - tmp_path: Path, -): - dest = import_auth_profile( - "work", - _write_auth( - tmp_path / "work.json", - access_token="work-token", - account_id="acct-work", - ), - ) - set_active_profile("work") - enable_auth_profile("work", enabled=False) - - assert load_codex_auth(dest) == ("work-token", "acct-work") - - -def test_rotation_randomly_selects_enabled_saved_profiles( - fake_home: Path, - tmp_path: Path, - monkeypatch: pytest.MonkeyPatch, - capsys, -): - import_auth_profile( - "alpha", - _write_auth( - tmp_path / "alpha.json", - access_token="alpha-token", - account_id="acct-alpha", - ), - ) - import_auth_profile( - "beta", - _write_auth( - tmp_path / "beta.json", - access_token="beta-token", - account_id="acct-beta", - ), - ) - set_active_profile("beta") - - assert main(["codex-lm", "rotation", "on"]) == 0 - capsys.readouterr() - - choices: list[list[str]] = [] - - def choose_beta(profiles): - choices.append([profile.name for profile in profiles]) - return next(profile for profile in profiles if profile.name == "beta") - - monkeypatch.setattr("dspy_codex_lm.auth.random.choice", choose_beta) - - assert load_codex_auth() == ("beta-token", "acct-beta") - assert load_codex_auth() == ("beta-token", "acct-beta") - assert choices == [["alpha", "beta"], ["alpha", "beta"]] - assert get_active_profile() == "beta" - rotation_state = fake_home / ".codex-lm" / "rotation.json" - assert "cursor" not in rotation_state.read_text(encoding="utf-8") - - def test_rotation_skips_disabled_profiles_and_all_disabled_fails( fake_home: Path, tmp_path: Path, @@ -398,17 +198,12 @@ def test_rotation_skips_disabled_profiles_and_all_disabled_fails( assert main(["codex-lm", "rotation", "on"]) == 0 capsys.readouterr() - choices: list[list[str]] = [] - def choose_only_enabled(profiles): - choices.append([profile.name for profile in profiles]) return profiles[0] monkeypatch.setattr("dspy_codex_lm.auth.random.choice", choose_only_enabled) assert load_codex_auth() == ("beta-token", "acct-beta") - assert load_codex_auth() == ("beta-token", "acct-beta") - assert choices == [["beta"], ["beta"]] enable_auth_profile("beta", enabled=False) @@ -466,67 +261,6 @@ def choose_alpha(profiles): assert random_choices == [["alpha", "beta"]] -def test_rotation_status_metadata_does_not_choose_request_profile( - fake_home: Path, - tmp_path: Path, - monkeypatch: pytest.MonkeyPatch, - capsys, -): - import_auth_profile( - "alpha", - _write_auth( - tmp_path / "alpha.json", - access_token="alpha-token", - account_id="acct-alpha", - ), - ) - import_auth_profile( - "beta", - _write_auth( - tmp_path / "beta.json", - access_token="beta-token", - account_id="acct-beta", - ), - ) - - assert main(["codex-lm", "rotation", "on"]) == 0 - capsys.readouterr() - - def fail_if_request_choice_used(profiles): - raise AssertionError(f"unexpected request selection from {profiles}") - - monkeypatch.setattr("dspy_codex_lm.auth.random.choice", fail_if_request_choice_used) - - metadata = auth_status_metadata() - - assert metadata["source"] == "rotation" - assert metadata["profile"] == "alpha" - - -def test_explicit_path_overrides_env_and_active( - fake_home: Path, - tmp_path: Path, - monkeypatch: pytest.MonkeyPatch, -): - import_auth_profile( - "active", - _write_auth( - tmp_path / "active.json", - access_token="active-token", - account_id="acct-active", - ), - ) - set_active_profile("active") - monkeypatch.setenv(CODEX_LM_AUTH_PROFILE_ENV, "active") - explicit = _write_auth( - tmp_path / "explicit.json", - access_token="explicit-token", - account_id="acct-explicit", - ) - - assert load_codex_auth(explicit) == ("explicit-token", "acct-explicit") - - def test_status_metadata_redacts_secrets(fake_home: Path, tmp_path: Path): import_auth_profile( "work", @@ -540,22 +274,51 @@ def test_status_metadata_redacts_secrets(fake_home: Path, tmp_path: Path): metadata = auth_status_metadata() - assert metadata["profile"] == "work" - assert metadata["source"] == "active" - assert metadata["access_token"] == "present" - assert metadata["refresh_token"] == "present" - assert metadata["id_token"] == "present" - assert metadata["account_id"] == "acct-v...ount" assert "secret-token" not in json.dumps(metadata) assert "person@example.com" not in json.dumps(metadata) + assert "acct-very-secret-account" not in json.dumps(metadata) + assert "id-secret" not in json.dumps(metadata) + assert metadata["access_token"] == "present" -@pytest.mark.parametrize("name", ["", ".", "..", "../work", "work/name", "work name"]) +@pytest.mark.parametrize("name", ["", "../work", "work/name"]) def test_invalid_profile_names_rejected(name: str): with pytest.raises(ValueError): validate_profile_name(name) -def test_clear_active_profile_is_idempotent(fake_home: Path): - clear_active_profile() - assert get_active_profile() is None +def test_long_lived_lm_refreshes_disabled_and_reenabled_accounts(fake_home, monkeypatch): + from conftest import build_stream_events + from dspy_codex_lm import CodexHTTPLM + + for name in ("alpha", "beta"): + import_auth_profile( + name, + _write_auth( + fake_home / f"{name}.json", + access_token=f"{name}-token", + account_id=f"acct-{name}", + ), + ) + assert main(["codex-lm", "rotation", "on"]) == 0 + now = 0.0 + monkeypatch.setattr("dspy_codex_lm.lm.monotonic", lambda: now) + monkeypatch.setattr( + "dspy_codex_lm.lm.random.choice", + lambda credentials: min(credentials, key=lambda credential: credential.account_id), + ) + + def transport(*, headers, api_key, **_): + account = headers["ChatGPT-Account-Id"] + assert api_key == account.removeprefix("acct-") + "-token" + return iter(build_stream_events(account)) + + monkeypatch.setattr("dspy_codex_lm.lm.litellm.responses", transport) + lm = CodexHTTPLM(model="gpt-5.3-codex", auth_config_refresh_seconds=60.0) + assert lm.forward(prompt="one", cache=False).output[0].content[0].text == "acct-alpha" + enable_auth_profile("alpha", enabled=False) + now = 60.0 + assert lm.forward(prompt="two", cache=False).output[0].content[0].text == "acct-beta" + enable_auth_profile("alpha", enabled=True) + now = 120.0 + assert lm.forward(prompt="three", cache=False).output[0].content[0].text == "acct-alpha" diff --git a/tests/codex_lm/test_auto_retry.py b/tests/codex_lm/test_auto_retry.py index 71eb7594..b0b98ca5 100644 --- a/tests/codex_lm/test_auto_retry.py +++ b/tests/codex_lm/test_auto_retry.py @@ -1,15 +1,4 @@ -"""CodexLM auto-retries transient stream failures. - -Codex streams drop or rate-limit frequently in long-running workloads -(24h+ optimize runs). Before this contract, a single dropped stream would -propagate up to GEPA and — depending on where the call lives — either -crash the proposer (losing an hour of inner-loop state) or zero out a -minibatch case. We now retry with exponential backoff so transient -errors are invisible above the LM layer. - -Persistent failures (all retries exhausted) still raise CodexStreamError -— the caller deserves to see the real failure mode, not a silent succeed. -""" +"""Retry recovery, exhaustion, and account failover.""" import copy import json @@ -53,67 +42,6 @@ def _failed_events( ] -def test_default_stream_attempts_match_upstream_codex(): - import dspy_codex_lm.lm as codex_lm - - assert codex_lm.DEFAULT_CODEX_STREAM_MAX_ATTEMPTS == 5 - - -def test_response_failed_retry_after_ms_is_attached_to_stream_error(lm): - from dspy_codex_lm.lm import _codex_stream_error_from_state - - state = lm._fresh_state() - lm._handle_event(_failed_events(retry_after_ms=1250)[0], state) - - error = _codex_stream_error_from_state(state) - - assert error is not None - assert error.retry_after_seconds == 1.25 - - -def test_retry_wait_prefers_server_requested_retry_after(monkeypatch): - from dspy_codex_lm.lm import _codex_retry_kwargs - - monkeypatch.setattr("dspy_codex_lm.lm.CODEX_STREAM_WAIT_MULTIPLIER", 0.0) - monkeypatch.setattr("dspy_codex_lm.lm.CODEX_STREAM_WAIT_MAX", 0.0) - wait = _codex_retry_kwargs()["wait"] - error = CodexStreamError("slow down", retry_after_seconds=3.5) - - class Outcome: - def exception(self): - return error - - class RetryState: - outcome = Outcome() - - assert wait(RetryState()) == 3.5 - - -def test_retries_on_transient_stream_failure_and_succeeds(lm, monkeypatch): - """One transient failure, then success → the caller sees the success - and the retry happens invisibly. - """ - # Override the conftest's test-default (max_attempts=1) to actually - # exercise retries. Wait knobs stay zero so the test is instant. - monkeypatch.setattr("dspy_codex_lm.lm.CODEX_STREAM_MAX_ATTEMPTS", 4) - - good_events = build_stream_events("ok", input_tokens=5, output_tokens=1) - call_sequence = [ - iter(copy.deepcopy(_failed_events())), - iter(copy.deepcopy(good_events)), - ] - - def fake(**_): - return call_sequence.pop(0) - - with mock.patch("dspy_codex_lm.lm.litellm.responses", side_effect=fake): - resp = lm.forward(prompt="hi") - - assert resp.output[0].content[0].text == "ok" - # Both the failing and succeeding stream were consumed - assert call_sequence == [] - - def test_retry_uses_alternate_rotation_profile_after_stream_stall(tmp_path, monkeypatch): from dspy_codex_lm import CodexHTTPLM as CodexLM from dspy_codex_lm.auth import import_auth_profile @@ -123,7 +51,9 @@ def test_retry_uses_alternate_rotation_profile_after_stream_stall(tmp_path, monk monkeypatch.setattr("dspy_codex_lm.lm.CODEX_STREAM_MAX_ATTEMPTS", 2) import_auth_profile( "alpha", - _write_auth(tmp_path / "alpha.json", access_token="alpha-token", account_id="acct-alpha"), + _write_auth( + tmp_path / "alpha.json", access_token="alpha-token", account_id="acct-alpha" + ), ) import_auth_profile( "beta", @@ -173,121 +103,9 @@ def fake(**_): return iter(copy.deepcopy(_failed_events(code="503", msg="upstream down"))) with mock.patch("dspy_codex_lm.lm.litellm.responses", side_effect=fake): - with pytest.raises(CodexStreamError, match="upstream down"): + with pytest.raises(CodexStreamError, match="upstream down") as caught: lm.forward(prompt="hi") - # max_attempts attempts were made in total - from dspy_codex_lm.lm import CODEX_STREAM_MAX_ATTEMPTS - - assert attempts["n"] == CODEX_STREAM_MAX_ATTEMPTS - - -async def test_aforward_retry_uses_alternate_rotation_profile_after_stream_stall( - tmp_path, - monkeypatch, -): - from dspy_codex_lm import CodexHTTPLM as CodexLM - from dspy_codex_lm.auth import import_auth_profile - from dspy_codex_lm.cli import main - - monkeypatch.setattr(Path, "home", classmethod(lambda cls: tmp_path)) - monkeypatch.setattr("dspy_codex_lm.lm.CODEX_STREAM_MAX_ATTEMPTS", 2) - import_auth_profile( - "alpha", - _write_auth(tmp_path / "alpha.json", access_token="alpha-token", account_id="acct-alpha"), - ) - import_auth_profile( - "beta", - _write_auth(tmp_path / "beta.json", access_token="beta-token", account_id="acct-beta"), - ) - assert main(["codex-lm", "rotation", "on"]) == 0 - - selected_accounts = iter(["acct-beta", "acct-beta"]) - - def choose_credentials(credentials): - credentials = tuple(credentials) - selected = next(selected_accounts) - return next( - credential for credential in credentials if credential.account_id == selected - ) - - seen_headers = [] - seen_api_keys = [] - good_events = build_stream_events("ok", input_tokens=5, output_tokens=1) - - async def good_stream(): - for event in copy.deepcopy(good_events): - yield event - - async def fake_aresponses(*, headers, api_key, **_): - seen_headers.append(headers["ChatGPT-Account-Id"]) - seen_api_keys.append(api_key) - if len(seen_headers) == 1: - raise CodexStreamError("Codex stream stalled") - return good_stream() - - monkeypatch.setattr("dspy_codex_lm.lm.random.choice", choose_credentials) - with mock.patch("dspy_codex_lm.lm.litellm.aresponses", side_effect=fake_aresponses): - response = await CodexLM(model="gpt-5.3-codex").aforward(prompt="hi") - - assert response.output[0].content[0].text == "ok" - assert seen_headers == ["acct-beta", "acct-alpha"] - assert seen_api_keys == ["beta-token", "alpha-token"] - - -async def test_aforward_retries_on_transient_failure(lm, monkeypatch): - """The async path mirrors the sync retry behaviour.""" - monkeypatch.setattr("dspy_codex_lm.lm.CODEX_STREAM_MAX_ATTEMPTS", 4) - - good_events = build_stream_events("ok", input_tokens=5, output_tokens=1) - - async def _make_fail_iter(): - for ev in _failed_events(): - yield ev - - async def _make_good_iter(): - for ev in copy.deepcopy(good_events): - yield ev - - call_sequence = [_make_fail_iter(), _make_good_iter()] - - async def fake(**_): - return call_sequence.pop(0) - - with mock.patch("dspy_codex_lm.lm.litellm.aresponses", side_effect=fake): - resp = await lm.aforward(prompt="hi") - - assert resp.output[0].content[0].text == "ok" - - -def test_pydantic_unexpected_value_warning_is_silenced(): - """``PydanticSerializationUnexpectedValue(... ResponseAPIUsage ...)`` - fires whenever a downstream pydantic model (DSPy's Prediction, a - user's logging wrapper, etc.) serializes a field declared as - ``ResponseAPIUsage`` whose actual value carries our Chat-Completions - aliases (``prompt_tokens`` / ``completion_tokens``). The aliases are - load-bearing for cost trackers, so rather than dropping them we - silence the resulting UserWarning at module load via - ``_install_pydantic_warning_filter``. - """ - import warnings - - from dspy_codex_lm.lm import _install_pydantic_warning_filter - - with warnings.catch_warnings(record=True) as caught: - warnings.resetwarnings() - _install_pydantic_warning_filter() - # Emit the exact warning shape the user saw from pydantic. - warnings.warn( - "Pydantic serializer warnings:\n" - " PydanticSerializationUnexpectedValue(Expected " - "`ResponseAPIUsage` - serialized value may not be as " - "expected [field_name='usage', input_value={...}])", - UserWarning, - stacklevel=2, - ) - - assert not caught, ( - f"expected ResponseAPIUsage serializer warnings to be silenced, " - f"got: {[str(w.message)[:80] for w in caught]}" - ) + assert caught.value.failure_kind == "failed" + assert caught.value.failure_code == "503" + assert attempts["n"] == 4 diff --git a/tests/codex_lm/test_build_request.py b/tests/codex_lm/test_build_request.py index 1e32a337..ff624bae 100644 --- a/tests/codex_lm/test_build_request.py +++ b/tests/codex_lm/test_build_request.py @@ -1,504 +1,40 @@ -import pytest +import os +from unittest import mock +from conftest import build_stream_events +from dspy_codex_lm import CodexHTTPLM -def test_build_request_from_prompt(lm): - request, headers = lm._build_request(prompt="hello world", messages=None, kwargs={}) - assert request["store"] is False - assert request["stream"] is True - assert request["instructions"] # non-empty - # Input was converted to list form (Responses API) - assert isinstance(request["input"], list) - assert request["input"][0]["role"] == "user" - assert request["input"][0]["content"][0]["type"] == "input_text" - assert request["input"][0]["content"][0]["text"] == "hello world" - - -def test_build_request_from_messages(lm): - request, _ = lm._build_request( - prompt=None, - messages=[{"role": "user", "content": "hi"}], - kwargs={}, - ) - assert request["input"][0]["content"][0]["text"] == "hi" - - -def test_build_request_headers_include_account_id(lm): - _, headers = lm._build_request(prompt="x", messages=None, kwargs={}) - assert headers["ChatGPT-Account-Id"] == "fake-account" - assert headers["originator"] == "opencode" - assert "session_id" in headers - - -@pytest.mark.parametrize( - "model", - ["gpt-5.6-sol", "gpt-5.6-terra", "gpt-5.6-luna"], -) -def test_gpt_5_6_models_use_responses_lite_header(model): - from dspy_codex_lm import CodexHTTPLM as CodexLM - - lm = CodexLM(model=model, access_token="fake", account_id="acct") - - request, headers = lm._build_request(prompt="x", messages=None, kwargs={}) - - assert headers["x-openai-internal-codex-responses-lite"] == "true" - assert { - "originator": headers["originator"], - "reasoning.context": request.get("reasoning", {}).get("context"), - "parallel_tool_calls": request.get("parallel_tool_calls"), - } == { - "originator": "codex_cli_rs", - "reasoning.context": "all_turns", - "parallel_tool_calls": False, - } - assert request["parallel_tool_calls"] is False - - -def test_gpt_5_6_request_does_not_mutate_constructor_reasoning(): - from dspy_codex_lm import CodexHTTPLM as CodexLM +def test_request_conversion_does_not_mutate_constructor_reasoning(): reasoning = {"effort": "high"} - lm = CodexLM( + lm = CodexHTTPLM( model="gpt-5.6-sol", access_token="fake", account_id="acct", reasoning=reasoning, ) - request, _ = lm._build_request(prompt="x", messages=None, kwargs={}) - assert request["reasoning"]["context"] == "all_turns" assert reasoning == {"effort": "high"} -def test_gpt_5_3_codex_does_not_use_responses_lite_header(lm): - request, headers = lm._build_request(prompt="x", messages=None, kwargs={}) - - assert "x-openai-internal-codex-responses-lite" not in headers - assert headers["originator"] == "opencode" - assert "reasoning" not in request - assert "parallel_tool_calls" not in request - - -def test_each_call_gets_unique_session_id(lm): - _, h1 = lm._build_request(prompt="x", messages=None, kwargs={}) - _, h2 = lm._build_request(prompt="x", messages=None, kwargs={}) - assert h1["session_id"] != h2["session_id"] - - -def test_build_request_drops_rollout_id_and_cache(lm): - request, _ = lm._build_request( - prompt="x", messages=None, kwargs={"rollout_id": "r1", "cache": True} - ) - assert "rollout_id" not in request - assert "cache" not in request - - -def test_build_request_canonicalizes_model_for_codex_backend(lm): - request, _ = lm._build_request(prompt="x", messages=None, kwargs={}) - assert request["model"] == "gpt-5.3-codex" - assert request["custom_llm_provider"] == "openai" - - -def test_build_request_canonicalizes_spark_model(): - from dspy_codex_lm import CodexHTTPLM as CodexLM - from litellm.utils import supports_native_streaming - - spark = CodexLM( - model="gpt-5.3-codex-spark", - access_token="fake", - account_id="acct", - ) - request, _ = spark._build_request(prompt="x", messages=None, kwargs={}) - assert request["model"] == "gpt-5.3-codex-spark" - assert request["stream"] is True - assert request["custom_llm_provider"] == "openai" - assert supports_native_streaming("gpt-5.3-codex-spark", "openai") is True - - -def test_custom_instructions_via_ctor(lm): - from dspy_codex_lm import CodexHTTPLM as CodexLM - - custom = CodexLM( - model="gpt-5.3-codex", - instructions="You are Ada.", +def test_proxy_environment_is_restored_after_request(monkeypatch): + monkeypatch.setenv("HTTPS_PROXY", "http://original:8080") + monkeypatch.delenv("HTTP_PROXY", raising=False) + lm = CodexHTTPLM( + model="gpt-5.5", + proxy_url="http://127.0.0.1:8898", access_token="fake", - account_id="acct", - ) - request, _ = custom._build_request(prompt="x", messages=None, kwargs={}) - assert request["instructions"] == "You are Ada." - - -def test_ctor_can_pin_auth_profile(tmp_path, monkeypatch): - import json - from pathlib import Path - - from dspy_codex_lm import CodexHTTPLM as CodexLM - from dspy_codex_lm.auth import import_auth_profile - - monkeypatch.setattr(Path, "home", classmethod(lambda cls: tmp_path)) - source = tmp_path / "source-auth.json" - source.write_text( - json.dumps( - { - "tokens": { - "access_token": "profile-token", - "account_id": "acct-profile", - } - } - ), - encoding="utf-8", + account_id="fake", ) - import_auth_profile("pro", source) - - custom = CodexLM(model="gpt-5.3-codex", auth_profile="pro") - - _, headers = custom._build_request(prompt="x", messages=None, kwargs={}) - assert headers["ChatGPT-Account-Id"] == "acct-profile" - - -def test_long_lived_lm_randomly_selects_accounts_per_request(tmp_path, monkeypatch): - import json - from pathlib import Path - from types import SimpleNamespace - - from dspy_codex_lm import CodexHTTPLM as CodexLM - from dspy_codex_lm.auth import import_auth_profile - from dspy_codex_lm.cli import main - - monkeypatch.setattr(Path, "home", classmethod(lambda cls: tmp_path)) - - def write_auth(path: Path, *, access_token: str, account_id: str) -> Path: - path.write_text( - json.dumps( - { - "tokens": { - "access_token": access_token, - "account_id": account_id, - } - } - ), - encoding="utf-8", - ) - return path - - import_auth_profile( - "alpha", - write_auth( - tmp_path / "alpha.json", - access_token="alpha-token", - account_id="acct-alpha", - ), - ) - import_auth_profile( - "beta", - write_auth( - tmp_path / "beta.json", - access_token="beta-token", - account_id="acct-beta", - ), - ) - assert main(["codex-lm", "rotation", "on"]) == 0 - - seen_headers = [] - seen_api_keys = [] - - def fake_responses(*, headers, api_key, **kwargs): - seen_headers.append(dict(headers)) - seen_api_keys.append(api_key) - return iter( - [ - SimpleNamespace(type="response.output_text.delta", delta="ok"), - SimpleNamespace( - type="response.completed", - response={ - "id": "resp", - "model": "gpt-5.3-codex", - "usage": { - "input_tokens": 1, - "output_tokens": 1, - "total_tokens": 2, - }, - }, - ), - ] - ) - monkeypatch.setattr("dspy_codex_lm.lm.litellm.responses", fake_responses) - - random_choices = [] - selected_accounts = iter(["acct-beta", "acct-alpha"]) - - def choose_credentials(credentials): - credentials = tuple(credentials) - random_choices.append([credential.account_id for credential in credentials]) - selected = next(selected_accounts) - return next( - credential for credential in credentials if credential.account_id == selected - ) - - monkeypatch.setattr("dspy_codex_lm.lm.random.choice", choose_credentials) - - lm = CodexLM(model="gpt-5.3-codex") - lm.forward(prompt="one", cache=False) - lm.forward(prompt="two", cache=False) - - assert [item["ChatGPT-Account-Id"] for item in seen_headers] == [ - "acct-beta", - "acct-alpha", - ] - assert seen_api_keys == ["beta-token", "alpha-token"] - assert random_choices == [["acct-alpha", "acct-beta"], ["acct-alpha", "acct-beta"]] - - -def test_long_lived_lm_resyncs_disabled_and_reenabled_profiles(tmp_path, monkeypatch): - import json - from pathlib import Path - - from dspy_codex_lm import CodexHTTPLM as CodexLM - from dspy_codex_lm.auth import enable_auth_profile, import_auth_profile - from dspy_codex_lm.cli import main - - monkeypatch.setattr(Path, "home", classmethod(lambda cls: tmp_path)) - now = 0.0 - monkeypatch.setattr("dspy_codex_lm.lm.monotonic", lambda: now) - - def write_auth(path: Path, *, access_token: str, account_id: str) -> Path: - path.write_text( - json.dumps( - { - "tokens": { - "access_token": access_token, - "account_id": account_id, - } - } - ), - encoding="utf-8", - ) - return path - - import_auth_profile( - "alpha", - write_auth( - tmp_path / "alpha.json", - access_token="alpha-token", - account_id="acct-alpha", - ), - ) - import_auth_profile( - "beta", - write_auth( - tmp_path / "beta.json", - access_token="beta-token", - account_id="acct-beta", - ), - ) - assert main(["codex-lm", "rotation", "on"]) == 0 - - selected_accounts = iter( - [ - "acct-alpha", - "acct-beta", - "acct-alpha", - "acct-beta", - "acct-beta", - "acct-alpha", - ] - ) - - def choose_credentials(credentials): - credentials = tuple(credentials) - selected = next(selected_accounts) - return next( - credential for credential in credentials if credential.account_id == selected - ) - - monkeypatch.setattr("dspy_codex_lm.lm.random.choice", choose_credentials) - - lm = CodexLM(model="gpt-5.3-codex", auth_config_refresh_seconds=60.0) - _, headers = lm._build_request(prompt="one", messages=None, kwargs={}) - enable_auth_profile("alpha", enabled=False) - _, stale_headers = lm._build_request(prompt="two", messages=None, kwargs={}) - _, still_stale_headers = lm._build_request( - prompt="three", - messages=None, - kwargs={}, - ) - now = 60.0 - _, refreshed_headers = lm._build_request(prompt="four", messages=None, kwargs={}) - enable_auth_profile("alpha", enabled=True) - _, stale_reenabled_headers = lm._build_request( - prompt="five", - messages=None, - kwargs={}, - ) - now = 120.0 - _, reenabled_headers = lm._build_request(prompt="six", messages=None, kwargs={}) - - assert [ - item["ChatGPT-Account-Id"] - for item in [ - headers, - stale_headers, - still_stale_headers, - refreshed_headers, - stale_reenabled_headers, - reenabled_headers, - ] - ] == [ - "acct-alpha", - "acct-beta", - "acct-alpha", - "acct-beta", - "acct-beta", - "acct-alpha", - ] - - -def test_long_lived_lm_randomly_selects_cached_snapshot_without_auth_reload( - tmp_path, - monkeypatch, -): - import json - from pathlib import Path - - from dspy_codex_lm import CodexHTTPLM as CodexLM - from dspy_codex_lm.auth import import_auth_profile - from dspy_codex_lm.cli import main - - monkeypatch.setattr(Path, "home", classmethod(lambda cls: tmp_path)) - monkeypatch.setattr("dspy_codex_lm.lm.monotonic", lambda: 0.0) - - def write_auth(path: Path, *, access_token: str, account_id: str) -> Path: - path.write_text( - json.dumps( - { - "tokens": { - "access_token": access_token, - "account_id": account_id, - } - } - ), - encoding="utf-8", - ) - return path - - import_auth_profile( - "alpha", - write_auth( - tmp_path / "alpha.json", - access_token="alpha-token", - account_id="acct-alpha", - ), - ) - import_auth_profile( - "beta", - write_auth( - tmp_path / "beta.json", - access_token="beta-token", - account_id="acct-beta", - ), - ) - assert main(["codex-lm", "rotation", "on"]) == 0 - - import dspy_codex_lm.lm as lm_module - - original_load_codex_auth = lm_module.load_codex_auth - load_calls = [] - - def counting_load_codex_auth(*args, **kwargs): - load_calls.append((args, kwargs)) - return original_load_codex_auth(*args, **kwargs) - - monkeypatch.setattr(lm_module, "load_codex_auth", counting_load_codex_auth) - - random_choices = [] - selected_accounts = iter(["acct-beta", "acct-alpha", "acct-beta"]) - - def choose_credentials(credentials): - credentials = tuple(credentials) - random_choices.append([credential.account_id for credential in credentials]) - selected = next(selected_accounts) - return next( - credential for credential in credentials if credential.account_id == selected - ) - - monkeypatch.setattr("dspy_codex_lm.lm.random.choice", choose_credentials) - - lm = CodexLM(model="gpt-5.3-codex", auth_config_refresh_seconds=60.0) - first, first_headers = lm._build_request(prompt="one", messages=None, kwargs={}) - second, second_headers = lm._build_request(prompt="two", messages=None, kwargs={}) - third, third_headers = lm._build_request(prompt="three", messages=None, kwargs={}) - - assert [ - item["ChatGPT-Account-Id"] - for item in [ - first_headers, - second_headers, - third_headers, - ] - ] == [ - "acct-beta", - "acct-alpha", - "acct-beta", - ] - assert [item["api_key"] for item in [first, second, third]] == [ - "beta-token", - "alpha-token", - "beta-token", - ] - assert random_choices == [ - ["acct-alpha", "acct-beta"], - ["acct-alpha", "acct-beta"], - ["acct-alpha", "acct-beta"], - ] - assert len(load_calls) == 2 - - -def test_pinned_access_token_bypasses_rotation_random_choice(tmp_path, monkeypatch): - import json - from pathlib import Path - - from dspy_codex_lm import CodexHTTPLM as CodexLM - from dspy_codex_lm.auth import import_auth_profile - from dspy_codex_lm.cli import main - - monkeypatch.setattr(Path, "home", classmethod(lambda cls: tmp_path)) - - def write_auth(path: Path, *, access_token: str, account_id: str) -> Path: - path.write_text( - json.dumps( - { - "tokens": { - "access_token": access_token, - "account_id": account_id, - } - } - ), - encoding="utf-8", - ) - return path - - import_auth_profile( - "alpha", - write_auth( - tmp_path / "alpha.json", - access_token="alpha-token", - account_id="acct-alpha", - ), - ) - assert main(["codex-lm", "rotation", "on"]) == 0 - - def fail_if_random_choice_used(credentials): - raise AssertionError(f"unexpected random rotation from {credentials}") - - monkeypatch.setattr("dspy_codex_lm.lm.random.choice", fail_if_random_choice_used) - - lm = CodexLM( - model="gpt-5.3-codex", - access_token="pinned-token", - account_id="acct-pinned", - ) - request, headers = lm._build_request(prompt="one", messages=None, kwargs={}) + def transport(**_): + assert os.environ["HTTPS_PROXY"] == "http://127.0.0.1:8898" + assert os.environ["HTTP_PROXY"] == "http://127.0.0.1:8898" + return iter(build_stream_events("ok")) - assert request["api_key"] == "pinned-token" - assert headers["ChatGPT-Account-Id"] == "acct-pinned" + with mock.patch("dspy_codex_lm.lm.litellm.responses", side_effect=transport): + response = lm.forward(prompt="hi", cache=False) + assert response.output[0].content[0].text == "ok" + assert os.environ["HTTPS_PROXY"] == "http://original:8080" + assert "HTTP_PROXY" not in os.environ diff --git a/tests/codex_lm/test_cache_hit_history.py b/tests/codex_lm/test_cache_hit_history.py deleted file mode 100644 index 33261e56..00000000 --- a/tests/codex_lm/test_cache_hit_history.py +++ /dev/null @@ -1,102 +0,0 @@ -"""Pins DSPy's cache-hit behavior as observed through CodexLM history. - -Background: ``dspy.clients.cache.Cache.get`` on a cache hit does -``response.usage = {}`` and ``response.cache_hit = True`` but does NOT clear -``response._hidden_params["response_cost"]``. As a result, ``BaseLM`` writes -an entry with empty usage and the original call's cost. Downstream cost -aggregators that sum ``entry["cost"]`` naively will double-count. - -These tests pin that contract so if DSPy changes the behaviour upstream -(e.g. zeros cost too, or keeps usage populated), we notice it here first. -""" - -from unittest import mock - -import pytest -from conftest import build_stream_events - - -def _patch_responses(events): - import copy - - def fake(**_): - return iter(copy.deepcopy(events)) - - return mock.patch("dspy_codex_lm.lm.litellm.responses", side_effect=fake) - - -def test_cache_miss_history_has_populated_usage(lm): - """Sanity: a cold call populates both input_tokens/output_tokens and - the Chat-Completions-shaped prompt_tokens/completion_tokens aliases. - """ - events = build_stream_events("hi", input_tokens=1000, output_tokens=50) - with _patch_responses(events): - lm(prompt="fresh") - - entry = lm.history[0] - usage = dict(entry["usage"]) - assert usage["input_tokens"] == 1000 - assert usage["output_tokens"] == 50 - # dspy-codex-lm aliases these so downstream Chat-Completions consumers work - assert usage["prompt_tokens"] == 1000 - assert usage["completion_tokens"] == 50 - assert entry["cost"] > 0 - - -def test_cache_hit_leaves_empty_usage_and_preserves_cost(lm): - """Upstream DSPy contract: cache hit zeros response.usage, keeps cost. - - If this test fails, DSPy has changed its Cache.get() semantics — update - predict-rlm's ``usage_since`` cache-hit detection accordingly. - """ - events = build_stream_events("hi", input_tokens=1000, output_tokens=50) - call_count = 0 - - def counting_fake(**_): - nonlocal call_count - call_count += 1 - import copy - - return iter(copy.deepcopy(events)) - - with mock.patch("dspy_codex_lm.lm.litellm.responses", side_effect=counting_fake): - lm(prompt="same prompt twice") - lm(prompt="same prompt twice") - - # Cache worked — only one underlying API call - assert call_count == 1 - assert len(lm.history) == 2 - - hit = lm.history[1] - hit_usage = dict(hit["usage"]) - # Cache hit: usage zeroed, cost preserved, cache_hit flag on response - assert hit_usage == {}, ( - f"DSPy cache-hit contract changed: usage={hit_usage!r} (expected empty)" - ) - assert hit["cost"] == pytest.approx(lm.history[0]["cost"]), ( - "DSPy cache-hit contract changed: cost was reset on cache hit " - "(previously preserved). predict-rlm's cache-hit cost discounting " - "may now be double-zeroing." - ) - assert getattr(hit["response"], "cache_hit", False) is True - - -def test_cache_hit_cost_would_double_count_without_filter(): - """Documents the naïve-sum failure mode: summing entry['cost'] across - fresh+cached entries reports 2x the real API spend even though only one - call was made. predict-rlm's ``usage_since`` filters this out. - """ - history = [ - {"usage": {"prompt_tokens": 1000, "completion_tokens": 50}, "cost": 0.001}, - {"usage": {}, "cost": 0.001}, # cache hit shape - ] - naive_total = sum(e["cost"] for e in history) - # Naïve sum reports 2x what was actually charged - assert naive_total == pytest.approx(0.002) - # Real API spend (only cache-MISS entries with populated usage) - real_total = sum( - e["cost"] - for e in history - if e["usage"].get("prompt_tokens", 0) or e["usage"].get("completion_tokens", 0) - ) - assert real_total == pytest.approx(0.001) diff --git a/tests/codex_lm/test_cli.py b/tests/codex_lm/test_cli.py index d3079aff..189f0be0 100644 --- a/tests/codex_lm/test_cli.py +++ b/tests/codex_lm/test_cli.py @@ -1,3 +1,4 @@ +import sys from pathlib import Path from types import SimpleNamespace from unittest import mock @@ -6,30 +7,18 @@ import pytest from dspy_codex_lm.auth import CODEX_LM_AUTH_PROFILE_ENV, load_codex_auth from dspy_codex_lm.cli import ( - CODEX_SUPPORTED_MODELS, - CodexLMUnsupportedModelError, - install_monkeypatch, - is_openai_family, main, - resolve_codex_model, - restore_monkeypatch, ) @pytest.fixture -def restore_lm(): - original_top = dspy.LM - try: - import dspy.clients.lm as mod +def restore_lm(monkeypatch): + import dspy.clients.lm as lm_module - original_inner = mod.LM - except Exception: - mod = None - original_inner = None - yield - dspy.LM = original_top - if mod is not None: - mod.LM = original_inner + monkeypatch.setattr(dspy, "LM", dspy.LM) + monkeypatch.setattr(lm_module, "LM", lm_module.LM) + monkeypatch.setattr(sys, "argv", sys.argv.copy()) + monkeypatch.setattr(sys, "path", sys.path.copy()) @pytest.fixture(autouse=True) @@ -62,8 +51,8 @@ def _write_auth( return path -@pytest.fixture -def fake_home(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Path: +@pytest.fixture(autouse=True) +def fake_home(tmp_path: Path, monkeypatch: pytest.MonkeyPatch, isolate_auth_home) -> Path: home = tmp_path / "home" home.mkdir() monkeypatch.setattr(Path, "home", classmethod(lambda cls: home)) @@ -72,213 +61,6 @@ def fake_home(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Path: return home -# ---- pattern matcher unit tests ---- - - -@pytest.mark.parametrize( - "model", - [ - "openai/gpt-4o", - "openai/o3", - "gpt-4o", - "gpt-4.1-mini", - "gpt-5", - "o1-mini", - "o3", - "o4-mini", - "chatgpt-4o-latest", - ], -) -def test_intercept_matches_openai_family(model: str): - assert is_openai_family(model) is True - - -@pytest.mark.parametrize( - "model", - [ - "anthropic/claude-3-5-sonnet", - "claude-3-opus", - "google/gemini-2.0-flash", - "gemini-pro", - "cohere/command-r", - "together_ai/meta-llama/Llama-3", - "ollama/llama3", - ], -) -def test_intercept_skips_other_providers(model: str): - assert is_openai_family(model) is False - - -# ---- resolver unit tests ---- - - -@pytest.mark.parametrize("slug", sorted(CODEX_SUPPORTED_MODELS)) -def test_resolve_returns_codex_slug_unchanged(slug: str): - assert resolve_codex_model(slug) == slug - - -def test_resolve_codex_model_strips_openai_prefix(): - assert resolve_codex_model("openai/gpt-5.3-codex") == "gpt-5.3-codex" - - -def test_resolve_codex_model_accepts_spark_slug(): - assert resolve_codex_model("openai/gpt-5.3-codex-spark") == "gpt-5.3-codex-spark" - -@pytest.mark.parametrize( - ("model", "expected"), - [ - ("gpt-5.6-sol", "gpt-5.6-sol"), - ("openai/gpt-5.6-sol", "gpt-5.6-sol"), - ("gpt-5.6-terra", "gpt-5.6-terra"), - ("openai/gpt-5.6-terra", "gpt-5.6-terra"), - ("gpt-5.6-luna", "gpt-5.6-luna"), - ("openai/gpt-5.6-luna", "gpt-5.6-luna"), - ], -) -def test_resolve_codex_model_accepts_gpt_5_6_family(model: str, expected: str): - assert resolve_codex_model(model) == expected - - - -@pytest.mark.parametrize( - "slug", - [ - "gpt-4o", - "openai/gpt-4o", - "gpt-4.1", - "gpt-5", - "gpt-5-mini", - "gpt-5.6", - "o3", - "o4-mini", - "openai/o3", - ], -) -def test_resolve_raises_for_unsupported(slug: str): - with pytest.raises(CodexLMUnsupportedModelError) as excinfo: - resolve_codex_model(slug) - msg = str(excinfo.value) - assert slug in msg - # Error message includes the supported set so users can self-serve - assert "gpt-5.3-codex" in msg - - -# ---- install / restore ---- - - -def test_install_and_restore(restore_lm): - original = dspy.LM - install_monkeypatch(verbose=False) - assert dspy.LM is not original - restore_monkeypatch(original, original) - assert dspy.LM is original - - -def test_install_intercepts_codex_slug(restore_lm): - install_monkeypatch(verbose=False) - lm = dspy.LM(model="openai/gpt-5.3-codex", api_key="sk-fake") - from dspy_codex_lm import CodexLM - - assert isinstance(lm, CodexLM) - assert lm.model == "openai/gpt-5.3-codex" - - -def test_lm_copy_does_not_trigger_intercept(restore_lm): - """``lm.copy()`` uses ``copy.deepcopy``, which calls - ``cls.__new__(cls)`` with NO model argument. Before the sentinel- - default fix, that kicked into the ``model="gpt-4o-mini"`` default - and raised ``CodexLMUnsupportedModelError``. Pin that deepcopy goes - through cleanly regardless of what the default is. - """ - install_monkeypatch(verbose=False) - lm = dspy.LM(model="anthropic/claude-opus-4-6", api_key="sk-fake") - # This triggers copy.deepcopy under the hood - lm_copy = lm.copy() - # No exception; copy behaves like the original - assert type(lm_copy) is type(lm) - assert lm_copy.model == lm.model - assert lm_copy.history == [] # DSPy's copy resets history - - -def test_codex_lm_copy_does_not_trigger_intercept(restore_lm): - """Same guarantee for an already-intercepted CodexLM instance.""" - install_monkeypatch(verbose=False) - lm = dspy.LM(model="openai/gpt-5.3-codex", api_key="sk-fake") - lm_copy = lm.copy() - assert type(lm_copy) is type(lm) - assert lm_copy.model == "openai/gpt-5.3-codex" - - -def test_intercepted_instance_tagged_with_source_model(restore_lm): - install_monkeypatch(verbose=False) - lm = dspy.LM(model="openai/gpt-5.3-codex", api_key="sk-fake") - assert lm.codex_intercepted_from == "openai/gpt-5.3-codex" - - -def test_passthrough_has_no_intercept_attr(restore_lm): - install_monkeypatch(verbose=False) - lm = dspy.LM(model="anthropic/claude-3-5-sonnet", api_key="sk-fake") - assert not hasattr(lm, "codex_intercepted_from") - - -def test_intercept_emits_info_log(restore_lm, caplog): - import logging as _logging - - caplog.set_level(_logging.INFO, logger="dspy_codex_lm.cli") - install_monkeypatch(verbose=False) - dspy.LM(model="openai/gpt-5.3-codex", api_key="sk-fake") - records = [r for r in caplog.records if r.name == "dspy_codex_lm.cli"] - assert len(records) == 1 - assert records[0].levelname == "INFO" - assert "intercept" in records[0].getMessage() - assert "openai/gpt-5.3-codex" in records[0].getMessage() - assert "gpt-5.3-codex" in records[0].getMessage() - - -def test_no_log_for_passthrough(restore_lm, caplog): - import logging as _logging - - caplog.set_level(_logging.INFO, logger="dspy_codex_lm.cli") - install_monkeypatch(verbose=False) - dspy.LM(model="anthropic/claude-3-5-sonnet", api_key="sk-fake") - records = [r for r in caplog.records if r.name == "dspy_codex_lm.cli"] - assert records == [] - - -def test_install_raises_on_unsupported_openai_model(restore_lm): - install_monkeypatch(verbose=False) - with pytest.raises(CodexLMUnsupportedModelError): - dspy.LM(model="openai/gpt-4o", api_key="sk-fake") - - -def test_install_passes_through_non_openai(restore_lm): - install_monkeypatch(verbose=False) - lm = dspy.LM(model="anthropic/claude-3-5-sonnet", api_key="sk-fake") - from dspy_codex_lm import CodexLM - - assert not isinstance(lm, CodexLM) - assert lm.model == "anthropic/claude-3-5-sonnet" - - -def test_install_drops_conflicting_kwargs(restore_lm): - install_monkeypatch(verbose=False) - # If we didn't strip api_key/api_base/model_type, CodexLM would get duplicate kwargs - lm = dspy.LM( - model="openai/gpt-5.3-codex", - api_key="sk-fake", - api_base="https://should-be-dropped", - model_type="chat", - ) - from dspy_codex_lm import CodexLM - - assert isinstance(lm, CodexLM) - # model_type should be "responses" (CodexLM's own value), not "chat" - assert lm.model_type == "responses" - - -# ---- main() end-to-end ---- - - def test_main_intercepts_in_child_script(tmp_path, restore_lm, capsys): script = _write_script( tmp_path, @@ -294,49 +76,6 @@ def test_main_intercepts_in_child_script(tmp_path, restore_lm, capsys): captured = capsys.readouterr() assert rc == 0 assert "OK" in captured.out - assert "intercept" in captured.err - - -def test_child_script_can_inspect_intercept_attr(tmp_path, restore_lm, capsys): - script = _write_script( - tmp_path, - """ -import dspy -lm = dspy.LM(model="openai/gpt-5.3-codex", api_key="sk-fake") -assert lm.codex_intercepted_from == "openai/gpt-5.3-codex" -print(f"came_from={lm.codex_intercepted_from}") -""", - ) - rc = main(["codex-lm", str(script)]) - captured = capsys.readouterr() - assert rc == 0 - assert "came_from=openai/gpt-5.3-codex" in captured.out - - -def test_child_script_can_configure_log_capture(tmp_path, restore_lm, capsys): - script = _write_script( - tmp_path, - """ -import logging -logs = [] - -class Grab(logging.Handler): - def emit(self, record): - logs.append(record.getMessage()) - -logging.getLogger("dspy_codex_lm.cli").addHandler(Grab()) -logging.getLogger("dspy_codex_lm.cli").setLevel(logging.INFO) - -import dspy -dspy.LM(model="openai/gpt-5.3-codex", api_key="sk-fake") -assert any("intercept" in m for m in logs), f"no intercept log captured: {logs}" -print(f"logged={logs[0]}") -""", - ) - rc = main(["codex-lm", str(script)]) - captured = capsys.readouterr() - assert rc == 0 - assert "logged=intercept" in captured.out def test_main_exits_3_on_unsupported_model(tmp_path, restore_lm, capsys): @@ -347,7 +86,6 @@ def test_main_exits_3_on_unsupported_model(tmp_path, restore_lm, capsys): rc = main(["codex-lm", str(script)]) captured = capsys.readouterr() assert rc == 3 - assert "cannot route" in captured.err assert "gpt-4o" in captured.err @@ -366,7 +104,6 @@ def test_main_passes_through_non_openai_in_child(tmp_path, restore_lm, capsys): captured = capsys.readouterr() assert rc == 0 assert "PASSTHROUGH" in captured.out - assert "intercept" not in captured.err def test_main_forwards_argv_to_child(tmp_path, restore_lm, capsys): @@ -386,130 +123,6 @@ def test_main_forwards_exit_code(tmp_path, restore_lm): assert main(["codex-lm", str(script)]) == 42 -def test_main_quiet_env_suppresses_banner(tmp_path, restore_lm, monkeypatch, capsys): - monkeypatch.setenv("CODEX_LM_QUIET", "1") - script = _write_script( - tmp_path, - "import dspy; dspy.LM(model='openai/gpt-5.3-codex')", - ) - main(["codex-lm", str(script)]) - captured = capsys.readouterr() - assert "intercept" not in captured.err - - -def test_main_missing_args_returns_usage(restore_lm, capsys): - assert main(["codex-lm"]) == 0 - captured = capsys.readouterr() - assert captured.err == "" - assert "usage: codex-lm" in captured.out - assert "codex-lm rotation" in captured.out - assert "random" in captured.out - assert "round robin" not in captured.out - assert "round-robin" not in captured.out - assert "Supported Codex models:" not in captured.out - - -def test_smoke_test_help_is_clean(capsys): - assert main(["codex-lm", "smoke-test", "--help"]) == 0 - captured = capsys.readouterr() - assert captured.out.strip() == ( - "usage: codex-lm smoke-test [PROFILE] [--model MODEL] [--prompt TEXT]" - ) - assert captured.err == "" - - -def test_main_usage_command_prints_redacted_summary( - fake_home: Path, - monkeypatch, - capsys, -): - def fake_fetch(): - return { - "rate_limit": { - "primary": { - "used": 1, - "limit": 4, - "remaining": 3, - "reset_at": "2026-05-10T12:00:00Z", - } - }, - "account_id": "acct-secret", - "email": "person@example.com", - "access_token": "secret-token", - } - - monkeypatch.setenv("CODEX_LM_ENABLE_LEGACY_AUTH_FALLBACK", "1") - monkeypatch.setattr("dspy_codex_lm.cli.fetch_codex_usage", fake_fetch) - - assert main(["codex-lm", "usage"]) == 0 - captured = capsys.readouterr() - assert captured.err == "" - lines = captured.out.splitlines() - assert lines[0] == "-" * 60 - assert lines[-1] == "-" * 60 - assert lines[1:-1] == [ - "Codex usage", - "rate_limit.primary: 3/4 remaining (75.0% remaining); resets 2026-05-10T12:00:00Z", - ] - assert "secret-token" not in captured.out - assert "acct-secret" not in captured.out - assert "person@example.com" not in captured.out - - -def test_main_usage_command_fetches_all_saved_profiles( - fake_home: Path, - tmp_path: Path, - monkeypatch, - capsys, -): - first = _write_auth( - tmp_path / "first.json", - access_token="first-token", - account_id="acct-first-secret", - ) - second = _write_auth( - tmp_path / "second.json", - access_token="second-token", - account_id="acct-second-secret", - ) - assert main(["codex-lm", "auth", "import", "work", "--from", str(first)]) == 0 - assert main(["codex-lm", "auth", "import", "personal", "--from", str(second)]) == 0 - capsys.readouterr() - - calls = [] - - def fake_fetch(*, auth_path=None): - calls.append(Path(auth_path)) - return { - "rate_limit": { - "primary_window": { - "used_percent": 25 if "work" in str(auth_path) else 50, - "limit_window_seconds": 18000, - "reset_after_seconds": 60, - } - }, - "access_token": "secret-token", - "account_id": "acct-secret", - "user": {"email": "person@example.com"}, - } - - monkeypatch.setattr("dspy_codex_lm.cli.fetch_codex_usage", fake_fetch) - - assert main(["codex-lm", "usage"]) == 0 - captured = capsys.readouterr() - - assert calls == [ - fake_home / ".codex-lm" / "auth" / "personal" / "auth.json", - fake_home / ".codex-lm" / "auth" / "work" / "auth.json", - ] - assert "personal:" in captured.out - assert "work (default):" in captured.out - assert captured.out.count("5h limit:") == 2 - assert "secret-token" not in captured.out - assert "acct-secret" not in captured.out - assert "person@example.com" not in captured.out - - def test_main_usage_skips_live_fetch_for_disabled_profiles( fake_home: Path, tmp_path: Path, @@ -554,355 +167,17 @@ def fake_fetch(*, auth_path=None): assert calls == [fake_home / ".codex-lm" / "auth" / "work" / "auth.json"] assert "personal (disabled):" in captured.out - assert " Disabled; live usage fetch skipped." in captured.out - assert "work (default):" in captured.out - assert "secret-token" not in captured.out - assert "acct-secret" not in captured.out - - -def test_main_usage_shows_rotation_on_for_saved_profiles( - fake_home: Path, - tmp_path: Path, - monkeypatch, - capsys, -): - first = _write_auth(tmp_path / "first.json") - second = _write_auth(tmp_path / "second.json") - assert main(["codex-lm", "auth", "import", "work", "--from", str(first)]) == 0 - assert main(["codex-lm", "auth", "import", "personal", "--from", str(second)]) == 0 - assert main(["codex-lm", "auth", "use", "work"]) == 0 - assert main(["codex-lm", "rotation", "on"]) == 0 - capsys.readouterr() - rotation_state = fake_home / ".codex-lm" / "rotation.json" - state_before = rotation_state.read_text(encoding="utf-8") - - def fake_fetch(*, auth_path=None): - return { - "rate_limit": { - "primary_window": { - "used_percent": 25 if "work" in str(auth_path) else 50, - "limit_window_seconds": 18000, - "reset_after_seconds": 60, - } - }, - "access_token": "secret-token", - "account_id": "acct-secret", - "user": {"email": "person@example.com"}, - } - - monkeypatch.setattr("dspy_codex_lm.cli.fetch_codex_usage", fake_fetch) - - assert main(["codex-lm", "usage"]) == 0 - captured = capsys.readouterr() - lines = captured.out.splitlines() - - assert lines[0] == "-" * 60 - assert lines[1:3] == ["Rotation: on (random)", ""] - assert "personal:" in captured.out assert "work (default):" in captured.out - assert captured.out.count("5h limit:") == 2 - assert rotation_state.read_text(encoding="utf-8") == state_before - assert "round robin" not in captured.out - assert "round-robin" not in captured.out assert "secret-token" not in captured.out assert "acct-secret" not in captured.out - assert "person@example.com" not in captured.out - - -def test_main_usage_color_can_be_forced_and_disabled( - fake_home: Path, - monkeypatch, - capsys, -): - def fake_fetch(): - return { - "rate_limit": { - "primary_window": { - "used_percent": 12, - "limit_window_seconds": 18000, - "reset_after_seconds": 300, - } - } - } - - monkeypatch.delenv("NO_COLOR", raising=False) - monkeypatch.setattr("dspy_codex_lm.cli.fetch_codex_usage", fake_fetch) - - assert main(["codex-lm", "--color=always", "usage"]) == 0 - captured = capsys.readouterr() - assert "\x1b[" in captured.out - assert "General usage limits:" in captured.out - - assert main(["codex-lm", "--no-color", "usage"]) == 0 - captured = capsys.readouterr() - assert "\x1b[" not in captured.out - assert "General usage limits:" in captured.out - - -def test_main_usage_respects_no_color_env(fake_home: Path, monkeypatch, capsys): - def fake_fetch(): - return { - "rate_limit": { - "primary_window": { - "used_percent": 12, - "limit_window_seconds": 18000, - "reset_after_seconds": 300, - } - } - } - - monkeypatch.setenv("NO_COLOR", "1") - monkeypatch.setattr("dspy_codex_lm.cli.fetch_codex_usage", fake_fetch) - - assert main(["codex-lm", "--color=always", "usage"]) == 0 - captured = capsys.readouterr() - assert "\x1b[" not in captured.out - assert "General usage limits:" in captured.out - - -def test_auth_import_list_use_status_remove(fake_home: Path, tmp_path: Path, capsys): - source = _write_auth( - tmp_path / "auth.json", - access_token="profile-token", - account_id="acct-profile-secret", - ) - - assert main(["codex-lm", "auth", "import", "work", "--from", str(source)]) == 0 - captured = capsys.readouterr() - assert captured.err == "" - assert "work" in captured.out - - assert main(["codex-lm", "auth", "list"]) == 0 - captured = capsys.readouterr() - assert captured.out.splitlines() == ["* work"] - - assert main(["codex-lm", "auth", "use", "work"]) == 0 - captured = capsys.readouterr() - assert captured.out.strip() == "Default auth profile: work" - - assert main(["codex-lm", "auth", "list"]) == 0 - captured = capsys.readouterr() - assert captured.out.splitlines() == ["* work"] - - assert main(["codex-lm", "auth", "status"]) == 0 - captured = capsys.readouterr() - assert "Active profile: work" in captured.out - assert "Selected profile: work (active profile)" in captured.out - assert "Access token: present" in captured.out - assert "Refresh token: present" in captured.out - assert "acct-p...cret" in captured.out - assert "profile-token" not in captured.out - assert "acct-profile-secret" not in captured.out - assert "person@example.com" not in captured.out - - assert main(["codex-lm", "auth", "remove", "work"]) == 0 - captured = capsys.readouterr() - assert captured.out.strip() == "Removed auth profile: work" - - assert main(["codex-lm", "auth", "list"]) == 0 - captured = capsys.readouterr() - assert captured.out.strip() == "No auth profiles saved." - - -def test_auth_disable_enable_and_list_marker(fake_home: Path, tmp_path: Path, capsys): - source = _write_auth( - tmp_path / "auth.json", - access_token="profile-token", - account_id="acct-profile-secret", - ) - - assert main(["codex-lm", "auth", "import", "work", "--from", str(source)]) == 0 - capsys.readouterr() - - assert main(["codex-lm", "auth", "disable", "work"]) == 0 - captured = capsys.readouterr() - assert captured.out.strip() == "Disabled auth profile: work" - assert "profile-token" not in captured.out - - assert main(["codex-lm", "auth", "list"]) == 0 - captured = capsys.readouterr() - assert captured.out.splitlines() == ["* work (disabled)"] - - assert main(["codex-lm", "auth", "enable", "work"]) == 0 - captured = capsys.readouterr() - assert captured.out.strip() == "Enabled auth profile: work" - - assert main(["codex-lm", "auth", "list"]) == 0 - captured = capsys.readouterr() - assert captured.out.splitlines() == ["* work"] - - -def test_auth_disable_missing_name_prints_usage(fake_home: Path, capsys): - assert main(["codex-lm", "auth", "disable"]) == 2 - captured = capsys.readouterr() - - assert "auth disable requires exactly one profile name" in captured.err - assert "usage: codex-lm auth disable NAME" in captured.err - - -def test_auth_use_and_default_alias_switch_load_codex_auth( - fake_home: Path, - tmp_path: Path, - capsys, -): - work = _write_auth( - tmp_path / "work.json", - access_token="work-token", - account_id="acct-work", - ) - personal = _write_auth( - tmp_path / "personal.json", - access_token="personal-token", - account_id="acct-personal", - ) - - assert main(["codex-lm", "auth", "import", "work", "--from", str(work)]) == 0 - assert main(["codex-lm", "auth", "import", "personal", "--from", str(personal)]) == 0 - capsys.readouterr() - - assert main(["codex-lm", "auth", "use", "work"]) == 0 - captured = capsys.readouterr() - assert captured.out.strip() == "Default auth profile: work" - assert load_codex_auth() == ("work-token", "acct-work") - - assert main(["codex-lm", "auth", "default", "personal"]) == 0 - captured = capsys.readouterr() - assert captured.out.strip() == "Default auth profile: personal" - assert load_codex_auth() == ("personal-token", "acct-personal") - - -def test_rotation_on_off_status( - fake_home: Path, - tmp_path: Path, - capsys, -): - first = _write_auth(tmp_path / "first.json") - assert main(["codex-lm", "auth", "import", "work", "--from", str(first)]) == 0 - capsys.readouterr() - - assert main(["codex-lm", "rotation"]) == 0 - captured = capsys.readouterr() - assert captured.err == "" - assert captured.out.strip() == "Rotation: off" - - assert main(["codex-lm", "rotation", "on"]) == 0 - captured = capsys.readouterr() - assert captured.err == "" - assert captured.out.strip() == "Rotation: on (random)" - rotation_state = fake_home / ".codex-lm" / "rotation.json" - assert "cursor" not in rotation_state.read_text(encoding="utf-8") - - assert main(["codex-lm", "rotation", "status"]) == 0 - captured = capsys.readouterr() - assert captured.err == "" - assert captured.out.strip() == "Rotation: on (random)" - assert "cursor" not in rotation_state.read_text(encoding="utf-8") - - assert main(["codex-lm", "rotation", "off"]) == 0 - captured = capsys.readouterr() - assert captured.err == "" - assert captured.out.strip() == "Rotation: off" - - -def test_rotation_on_without_profiles_errors( - fake_home: Path, - capsys, -): - assert main(["codex-lm", "rotation", "on"]) == 2 - captured = capsys.readouterr() - assert captured.out == "" - assert "no auth profiles saved" in captured.err - assert "secret-token" not in captured.err - assert "person@example.com" not in captured.err - - -def test_auth_list_shows_rotation_on_and_keeps_default_marker( - fake_home: Path, - tmp_path: Path, - capsys, -): - first = _write_auth(tmp_path / "first.json") - second = _write_auth(tmp_path / "second.json") - assert main(["codex-lm", "auth", "import", "work", "--from", str(first)]) == 0 - assert main(["codex-lm", "auth", "import", "personal", "--from", str(second)]) == 0 - assert main(["codex-lm", "auth", "use", "work"]) == 0 - capsys.readouterr() - - assert main(["codex-lm", "auth", "list"]) == 0 - captured = capsys.readouterr() - assert captured.out.splitlines() == [" personal", "* work"] - assert "Rotation:" not in captured.out - - assert main(["codex-lm", "rotation", "on"]) == 0 - capsys.readouterr() - rotation_state = fake_home / ".codex-lm" / "rotation.json" - state_before = rotation_state.read_text(encoding="utf-8") - - assert main(["codex-lm", "auth", "list"]) == 0 - captured = capsys.readouterr() - - assert captured.out.splitlines() == [ - "Rotation: on (random)", - "", - " personal", - "* work", - ] - assert rotation_state.read_text(encoding="utf-8") == state_before - assert "round robin" not in captured.out - assert "round-robin" not in captured.out - - -def test_auth_list_and_status_color_can_be_forced( - fake_home: Path, - tmp_path: Path, - monkeypatch, - capsys, -): - monkeypatch.delenv("NO_COLOR", raising=False) - source = _write_auth(tmp_path / "auth.json") - - assert main(["codex-lm", "auth", "import", "work", "--from", str(source)]) == 0 - assert main(["codex-lm", "auth", "use", "work"]) == 0 - capsys.readouterr() - - assert main(["codex-lm", "--color=always", "auth", "list"]) == 0 - captured = capsys.readouterr() - assert "\x1b[" in captured.out - assert "* " in captured.out - assert "work" in captured.out - - assert main(["codex-lm", "--color=always", "auth", "status"]) == 0 - captured = capsys.readouterr() - assert "\x1b[" in captured.out - assert "Codex auth status" in captured.out - assert "secret-token" not in captured.out - - -def test_auth_status_shows_env_override( - fake_home: Path, - tmp_path: Path, - monkeypatch, - capsys, -): - source = _write_auth(tmp_path / "auth.json", account_id="acct-env-secret") - assert main(["codex-lm", "auth", "import", "envprofile", "--from", str(source)]) == 0 - capsys.readouterr() - monkeypatch.setenv(CODEX_LM_AUTH_PROFILE_ENV, "envprofile") - - assert main(["codex-lm", "auth", "status"]) == 0 - captured = capsys.readouterr() - - assert "Selected profile: envprofile (CODEX_LM_AUTH_PROFILE)" in captured.out - assert "acct-e...cret" in captured.out - assert "secret-token" not in captured.out - assert "person@example.com" not in captured.out def test_auth_login_uses_isolated_codex_home(fake_home: Path, monkeypatch, capsys): seen = {} + existing = _write_auth(fake_home / ".codex" / "auth.json") + before = existing.read_bytes() def fake_run(cmd, *, env): - seen["cmd"] = cmd seen["codex_home"] = env["CODEX_HOME"] _write_auth( Path(env["CODEX_HOME"]) / "auth.json", @@ -914,80 +189,18 @@ def fake_run(cmd, *, env): monkeypatch.setattr("dspy_codex_lm.cli.subprocess.run", fake_run) assert main(["codex-lm", "auth", "login", "login-profile"]) == 0 - captured = capsys.readouterr() - - assert seen["cmd"] == ["codex", "login", "--device-auth"] - assert Path(seen["codex_home"]).name.startswith("codex-lm-auth-") - assert str(fake_home) not in seen["codex_home"] - assert "login-profile" in captured.out + capsys.readouterr() + assert existing.read_bytes() == before + assert not Path(seen["codex_home"]).exists() + assert load_codex_auth(profile="login-profile") == ("login-token", "acct-login-secret") assert main(["codex-lm", "auth", "use", "login-profile"]) == 0 capsys.readouterr() assert main(["codex-lm", "auth", "status"]) == 0 captured = capsys.readouterr() - assert "acct-l...cret" in captured.out - assert "login-token" not in captured.out - - -def test_auth_login_missing_name_prints_usage(fake_home: Path, capsys): - assert main(["codex-lm", "auth", "login"]) == 2 - captured = capsys.readouterr() - - assert "auth login requires a profile name" in captured.err - assert "usage: codex-lm auth login NAME [--device-auth]" in captured.err - - -def test_auth_login_accepts_display_name_and_preserves_it( - fake_home: Path, - monkeypatch, - capsys, -): - seen = {} - - def fake_run(cmd, *, env): - seen["cmd"] = cmd - _write_auth( - Path(env["CODEX_HOME"]) / "auth.json", - access_token="login-token", - account_id="acct-login-secret", - ) - return SimpleNamespace(returncode=0) - - monkeypatch.setattr("dspy_codex_lm.cli.subprocess.run", fake_run) - - display_name = "gabriel@trampoline.ai" - slug = "gabriel-trampoline.ai" - - assert main(["codex-lm", "auth", "login", display_name]) == 0 - captured = capsys.readouterr() - - assert seen["cmd"] == ["codex", "login", "--device-auth"] - assert display_name in captured.out - profile_dir = fake_home / ".codex-lm" / "auth" / slug - assert (profile_dir / "auth.json").is_file() - assert (profile_dir / "profile.json").read_text(encoding="utf-8") == ( - '{\n "name": "gabriel@trampoline.ai",\n "slug": "gabriel-trampoline.ai"\n}\n' - ) - - assert main(["codex-lm", "auth", "list"]) == 0 - captured = capsys.readouterr() - assert captured.out.splitlines() == [f"* {display_name}"] - - assert main(["codex-lm", "auth", "use", display_name]) == 0 - captured = capsys.readouterr() - assert captured.out.strip() == f"Default auth profile: {display_name}" - - assert main(["codex-lm", "auth", "status"]) == 0 - captured = capsys.readouterr() - assert f"Active profile: {display_name}" in captured.out - assert f"Selected profile: {display_name} (active profile)" in captured.out assert "login-token" not in captured.out assert "acct-login-secret" not in captured.out - - assert main(["codex-lm", "auth", "remove", display_name]) == 0 - captured = capsys.readouterr() - assert captured.out.strip() == f"Removed auth profile: {display_name}" - assert not profile_dir.exists() + assert "secret-refresh" not in captured.out def test_auth_login_returns_codex_exit_code(fake_home: Path, monkeypatch, capsys): @@ -999,143 +212,4 @@ def fake_run(cmd, *, env): assert main(["codex-lm", "auth", "login", "work", "--device-auth"]) == 17 captured = capsys.readouterr() assert captured.out == "" - - -def test_auth_login_sets_default_when_none(fake_home: Path, monkeypatch, capsys): - def fake_run(cmd, *, env): - _write_auth( - Path(env["CODEX_HOME"]) / "auth.json", - access_token="login-token", - account_id="acct-login-secret", - ) - return SimpleNamespace(returncode=0) - - monkeypatch.setattr("dspy_codex_lm.cli.subprocess.run", fake_run) - - assert main(["codex-lm", "auth", "login", "first-login"]) == 0 - capsys.readouterr() - - assert load_codex_auth() == ("login-token", "acct-login-secret") - - -def test_smoke_test_checks_each_saved_profile( - fake_home: Path, - tmp_path: Path, - monkeypatch, - capsys, -): - first = _write_auth( - tmp_path / "first.json", - access_token="first-token", - account_id="acct-first-secret", - ) - second = _write_auth( - tmp_path / "second.json", - access_token="second-token", - account_id="acct-second-secret", - ) - assert main(["codex-lm", "auth", "import", "work", "--from", str(first)]) == 0 - assert main(["codex-lm", "auth", "import", "personal", "--from", str(second)]) == 0 - capsys.readouterr() - - calls = [] - - class FakeCodexLM: - def __init__(self, *, model, auth_profile=None): - calls.append(("init", model, auth_profile)) - - def forward(self, *, prompt, cache): - calls.append(("forward", prompt, cache)) - return SimpleNamespace() - - monkeypatch.setattr("dspy_codex_lm.cli.CodexLM", FakeCodexLM) - - assert main(["codex-lm", "smoke-test", "--model", "gpt-5.4-mini"]) == 0 - captured = capsys.readouterr() - - assert calls == [ - ("init", "gpt-5.4-mini", "personal"), - ("forward", "Reply with OK.", False), - ("init", "gpt-5.4-mini", "work"), - ("forward", "Reply with OK.", False), - ] - assert captured.err == "" - assert "personal: ok" in captured.out - assert "work: ok" in captured.out - - -def test_smoke_test_checks_only_requested_profile( - fake_home: Path, - tmp_path: Path, - monkeypatch, - capsys, -): - first = _write_auth(tmp_path / "first.json") - second = _write_auth(tmp_path / "second.json") - assert main(["codex-lm", "auth", "import", "work", "--from", str(first)]) == 0 - assert main(["codex-lm", "auth", "import", "personal", "--from", str(second)]) == 0 - capsys.readouterr() - - calls = [] - - class FakeCodexLM: - def __init__(self, *, model, auth_profile=None): - calls.append(("init", model, auth_profile)) - - def forward(self, *, prompt, cache): - calls.append(("forward", prompt, cache)) - return SimpleNamespace() - - monkeypatch.setattr("dspy_codex_lm.cli.CodexLM", FakeCodexLM) - - assert main(["codex-lm", "smoke-test", "work", "--prompt", "ping"]) == 0 - captured = capsys.readouterr() - - assert calls == [ - ("init", "gpt-5.5", "work"), - ("forward", "ping", False), - ] - assert "work: ok" in captured.out - assert "personal" not in captured.out - - -def test_auth_rejects_invalid_profile_name(fake_home: Path, capsys): - assert main(["codex-lm", "auth", "use", "../bad"]) == 2 - captured = capsys.readouterr() - assert "invalid auth profile name" in captured.err - assert "usage: codex-lm auth use NAME" in captured.err - - -@pytest.mark.parametrize( - ("argv", "usage"), - [ - ( - ["codex-lm", "auth", "import", "../bad"], - "usage: codex-lm auth import NAME [--from PATH]", - ), - ( - ["codex-lm", "auth", "login", "../bad"], - "usage: codex-lm auth login NAME [--device-auth]", - ), - ( - ["codex-lm", "auth", "remove", "../bad"], - "usage: codex-lm auth remove NAME", - ), - ], -) -def test_auth_invalid_profile_input_prints_command_usage( - fake_home: Path, - monkeypatch, - capsys, - argv, - usage, -): - run = mock.Mock() - monkeypatch.setattr("dspy_codex_lm.cli.subprocess.run", run) - - assert main(argv) == 2 - captured = capsys.readouterr() - - assert "invalid auth profile name" in captured.err - assert usage in captured.err - run.assert_not_called() + assert not (fake_home / ".codex-lm" / "auth" / "work").exists() diff --git a/tests/codex_lm/test_concurrent_cache.py b/tests/codex_lm/test_concurrent_cache.py index b0bae082..a79ed60f 100644 --- a/tests/codex_lm/test_concurrent_cache.py +++ b/tests/codex_lm/test_concurrent_cache.py @@ -1,51 +1,47 @@ import asyncio -import time from unittest import mock from conftest import build_stream_events -async def _delayed_async_iter(items, delay): - await asyncio.sleep(delay) - for item in items: - yield item - - -def _patch_slow_aresponses(events, delay): - async def fake(**_): - # Each call gets its own async iterator so concurrent iteration works - return _delayed_async_iter(events, delay) - - return mock.patch("dspy_codex_lm.lm.litellm.aresponses", side_effect=fake) - - -async def test_concurrent_cache_hits_fast(lm): - events = build_stream_events("answer", input_tokens=10, output_tokens=3) - with _patch_slow_aresponses(events, delay=0.1): - # Warm the cache with one call - await lm.aforward(prompt="same question") - - # Now patch with a fresh slow responder, but all these calls should hit cache - with _patch_slow_aresponses(events, delay=10.0): # if cache misses we'd hang - t0 = time.monotonic() - results = await asyncio.gather(*[lm.aforward(prompt="same question") for _ in range(5)]) - dt = time.monotonic() - t0 - - assert len(results) == 5 - assert all(r.output[0].content[0].text == "answer" for r in results) - # 5 cache hits should complete far under the 10s delay - assert dt < 0.5, f"cache hits took {dt:.3f}s, expected << 10s" - - -async def test_concurrent_misses_run_in_parallel(lm): - """5 distinct prompts = 5 misses. Should run concurrently, not serially.""" - events = build_stream_events("ok", input_tokens=10, output_tokens=1) - delay = 0.1 - with _patch_slow_aresponses(events, delay=delay): - t0 = time.monotonic() - await asyncio.gather(*[lm.aforward(prompt=f"question {i}") for i in range(5)]) - dt = time.monotonic() - t0 - - # Serialized would be ~5*delay = 0.5s; parallel should be ~delay = 0.1s. - # Allow some headroom for scheduling overhead. - assert dt < delay * 2.5, f"concurrent misses took {dt:.3f}s, expected ~{delay:.2f}s" +async def test_concurrent_requests_keep_results_and_cache_entries_isolated(lm): + entered = 0 + both_entered = asyncio.Event() + streams = iter( + [ + build_stream_events("alpha", input_tokens=10, output_tokens=1), + build_stream_events("beta", input_tokens=20, output_tokens=2), + ] + ) + + async def fake_aresponses(**_): + nonlocal entered + events = next(streams) + entered += 1 + if entered == 2: + both_entered.set() + + async def stream(): + await both_entered.wait() + for event in events: + yield event + await asyncio.sleep(0) + + return stream() + + prompts = ["first", "second"] + with mock.patch("dspy_codex_lm.lm.litellm.aresponses", side_effect=fake_aresponses): + fresh = await asyncio.wait_for( + asyncio.gather(*(lm.aforward(prompt=prompt) for prompt in prompts)), + timeout=3, + ) + assert {response.output[0].content[0].text for response in fresh} == {"alpha", "beta"} + assert {response.usage.input_tokens for response in fresh} == {10, 20} + expected = [response.output[0].content[0].text for response in fresh] + + with mock.patch( + "dspy_codex_lm.lm.litellm.aresponses", + side_effect=AssertionError("cache hit reached transport"), + ): + cached = await asyncio.gather(*(lm.aforward(prompt=prompt) for prompt in prompts)) + assert [response.output[0].content[0].text for response in cached] == expected diff --git a/tests/codex_lm/test_cost.py b/tests/codex_lm/test_cost.py deleted file mode 100644 index 50f81d8f..00000000 --- a/tests/codex_lm/test_cost.py +++ /dev/null @@ -1,78 +0,0 @@ -from dspy_codex_lm.cost import compute_cost - - -def test_gpt_5_3_codex_normal_pricing(): - # Published rates: input 1.75e-6, output 1.4e-5 - usage = { - "input_tokens": 100, - "output_tokens": 10, - "input_tokens_details": {"cached_tokens": 0}, - } - cost = compute_cost("gpt-5.3-codex", usage) - expected = (100 * 1.75e-6) + (10 * 1.4e-5) - assert abs(cost - expected) < 1e-12 - - -def test_strips_openai_prefix_from_model_slug(): - # LM stores model as "openai/gpt-5.3-codex" — compute_cost should strip it - usage = {"input_tokens": 100, "output_tokens": 10} - direct = compute_cost("gpt-5.3-codex", usage) - prefixed = compute_cost("openai/gpt-5.3-codex", usage) - assert direct == prefixed > 0 - - -def test_cached_tokens_priced_at_cache_rate(): - # 100 total input, 80 cached → 20 at normal rate, 80 at cache rate - usage = { - "input_tokens": 100, - "output_tokens": 0, - "input_tokens_details": {"cached_tokens": 80}, - } - cost = compute_cost("gpt-5.3-codex", usage) - # gpt-5.3-codex: input 1.75e-6, cache 1.75e-7 - expected = (20 * 1.75e-6) + (80 * 1.75e-7) - assert abs(cost - expected) < 1e-12 - - -def test_unknown_model_returns_zero(): - usage = {"input_tokens": 1000, "output_tokens": 500} - assert compute_cost("does-not-exist", usage) == 0.0 - - -def test_empty_usage_returns_zero(): - assert compute_cost("gpt-5.3-codex", {}) == 0.0 - - -def test_missing_details_key_does_not_crash(): - usage = {"input_tokens": 100, "output_tokens": 10} - # No "input_tokens_details" key — should treat cached as 0 - cost = compute_cost("gpt-5.3-codex", usage) - expected = (100 * 1.75e-6) + (10 * 1.4e-5) - assert abs(cost - expected) < 1e-12 - - -def test_gpt_5_4_has_its_own_rates(): - # Make sure we look up the actual slug (5.4), not collapse to 5.3 - usage = {"input_tokens": 100, "output_tokens": 10} - c_5_3 = compute_cost("gpt-5.3-codex", usage) - c_5_4 = compute_cost("gpt-5.4", usage) - # gpt-5.4 is more expensive than gpt-5.3-codex - assert c_5_4 > c_5_3 - - -def test_mercury_2_pricing_from_inception_rate_card(): - usage = { - "input_tokens": 1_000_000, - "output_tokens": 1_000_000, - "input_tokens_details": {"cached_tokens": 100_000}, - } - cost = compute_cost("openai/mercury-2", usage) - expected = (900_000 * 0.25e-6) + (100_000 * 0.025e-6) + (1_000_000 * 0.75e-6) - assert abs(cost - expected) < 1e-12 - - -def test_mercury_edit_2_pricing_from_inception_rate_card(): - usage = {"input_tokens": 1_000_000, "output_tokens": 1_000_000} - cost = compute_cost("mercury-edit-2", usage) - expected = 0.25 + 0.75 - assert abs(cost - expected) < 1e-12 diff --git a/tests/codex_lm/test_events.py b/tests/codex_lm/test_events.py deleted file mode 100644 index 26e2da7f..00000000 --- a/tests/codex_lm/test_events.py +++ /dev/null @@ -1,46 +0,0 @@ -from conftest import make_completed, make_text_delta - - -def test_text_delta_accumulates(lm): - state = lm._fresh_state() - lm._handle_event(make_text_delta("he"), state) - lm._handle_event(make_text_delta("llo"), state) - assert "".join(state["text_parts"]) == "hello" - - -def test_completed_captures_usage_and_ids(lm): - state = lm._fresh_state() - lm._handle_event( - make_completed( - input_tokens=20, - output_tokens=8, - response_id="resp-xyz", - model="gpt-5.4", - ), - state, - ) - assert state["response_id"] == "resp-xyz" - assert state["model_name"] == "gpt-5.4" - assert state["usage_raw"] is not None - - -def test_unknown_event_ignored(lm): - from types import SimpleNamespace - - state = lm._fresh_state() - original = dict(state) - lm._handle_event(SimpleNamespace(type="some.unrelated.event"), state) - assert state == original - - -def test_none_event_safe(lm): - state = lm._fresh_state() - original = dict(state) - lm._handle_event(None, state) - assert state == original - - -def test_empty_delta_not_appended(lm): - state = lm._fresh_state() - lm._handle_event(make_text_delta(""), state) - assert state["text_parts"] == [] diff --git a/tests/codex_lm/test_forward.py b/tests/codex_lm/test_forward.py index 9b62ea87..2a530c1c 100644 --- a/tests/codex_lm/test_forward.py +++ b/tests/codex_lm/test_forward.py @@ -1,82 +1,72 @@ +import copy +from types import SimpleNamespace from unittest import mock import dspy -from conftest import build_stream_events - - -def _patch_responses(events, call_count=None): - """side_effect that returns a fresh iterator on every call.""" - - def fake(**_): - if call_count is not None: - call_count["n"] += 1 - return iter(events) - - return mock.patch("dspy_codex_lm.lm.litellm.responses", side_effect=fake) - - -def test_forward_single_call(lm): - events = build_stream_events("4", input_tokens=100, output_tokens=1) - with _patch_responses(events): - resp = lm.forward(prompt="What is 2+2?") - assert resp.output[0].content[0].text == "4" - assert resp.usage.input_tokens == 100 - assert resp.usage.cost > 0 - - -def test_forward_disables_litellm_stream_logging(lm): - class StreamWithLitellmLogging: - def __init__(self, events): - self._events = iter(events) - - def __iter__(self): - return self - - def __next__(self): - event = next(self._events) - if event.type == "response.completed": - self._handle_logging_completed_response() - return event - - def _handle_logging_completed_response(self): - raise AssertionError("LiteLLM stream logging should be disabled") - - events = build_stream_events("ok", input_tokens=10, output_tokens=2) - - def fake(**_): - return StreamWithLitellmLogging(events) - - with mock.patch("dspy_codex_lm.lm.litellm.responses", side_effect=fake): - resp = lm.forward(prompt="ping", cache=False) - - assert resp.output[0].content[0].text == "ok" - - -def test_forward_via_dspy_predict(lm): - text = "[[ ## answer ## ]]\n4\n[[ ## completed ## ]]" - events = build_stream_events(text, input_tokens=50, output_tokens=10) - with _patch_responses(events): - dspy.configure(lm=lm) - result = dspy.Predict("question -> answer")(question="What is 2+2?") - assert result.answer.strip() == "4" - - -def test_forward_cost_in_history(lm): - text = "[[ ## answer ## ]]\n4\n[[ ## completed ## ]]" - events = build_stream_events(text, input_tokens=100, output_tokens=10) - with _patch_responses(events): - dspy.configure(lm=lm) - dspy.Predict("question -> answer")(question="What is 2+2?") - history_cost = dspy.settings.lm.history[-1]["cost"] - # 100 * 1.75e-6 + 10 * 1.4e-5 = 0.000315 - assert abs(history_cost - 0.000315) < 1e-9 - - -def test_forward_second_call_hits_cache(lm): - events = build_stream_events("first", input_tokens=10, output_tokens=3) - count = {"n": 0} - with _patch_responses(events, call_count=count): - r1 = lm.forward(prompt="same question") - r2 = lm.forward(prompt="same question") - assert count["n"] == 1 # second call hit cache - assert r1.output[0].content[0].text == r2.output[0].content[0].text == "first" +import pytest +from conftest import build_stream_events, make_completed, make_text_delta +from dspy.utils.usage_tracker import UsageTracker + + +def test_stream_assembly_accounts_for_cached_input_and_output(lm): + events = [ + make_text_delta("hel"), + SimpleNamespace(type="response.in_progress"), + make_text_delta("lo"), + make_completed(input_tokens=100, cached_tokens=80, output_tokens=10), + ] + with mock.patch("dspy_codex_lm.lm.litellm.responses", return_value=iter(events)): + response = lm.forward(prompt="hi", cache=False) + + assert response.output[0].content[0].text == "hello" + assert response.usage.input_tokens == 100 + assert response.usage.output_tokens == 10 + assert response.usage.cost == pytest.approx(20 * 1.75e-6 + 80 * 1.75e-7 + 10 * 1.4e-5) + + +def test_cached_calls_do_not_double_count_usage_or_mutate_fresh_history(lm): + events = build_stream_events("hi", input_tokens=1000, output_tokens=50) + tracker = UsageTracker() + with dspy.settings.context(usage_tracker=tracker): + with mock.patch( + "dspy_codex_lm.lm.litellm.responses", + side_effect=lambda **_: iter(copy.deepcopy(events)), + ) as transport: + assert lm(prompt="same")[0]["text"] == "hi" + assert lm(prompt="same")[0]["text"] == "hi" + + assert transport.call_count == 1 + totals = tracker.get_total_tokens()[lm.model] + assert totals["prompt_tokens"] == 1000 + assert totals["completion_tokens"] == 50 + fresh, cached = lm.history + assert fresh["usage"]["prompt_tokens"] == 1000 + assert fresh["usage"]["completion_tokens"] == 50 + assert fresh["cost"] == pytest.approx(1000 * 1.75e-6 + 50 * 1.4e-5) + assert dict(cached["usage"]) == {} + + +async def test_async_prediction_exposes_parsed_answer_and_billable_usage(lm): + events = build_stream_events( + "[[ ## answer ## ]]\n42\n[[ ## completed ## ]]", + input_tokens=500, + output_tokens=20, + ) + + async def fake_aresponses(**_): + async def stream(): + for event in copy.deepcopy(events): + yield event + + return stream() + + with dspy.settings.context(lm=lm, track_usage=True): + with mock.patch("dspy_codex_lm.lm.litellm.aresponses", side_effect=fake_aresponses): + prediction = await dspy.Predict("question -> answer").acall( + question="six times seven?" + ) + + assert prediction.answer.strip() == "42" + usage = prediction.get_lm_usage()[lm.model] + assert usage["prompt_tokens"] == 500 + assert usage["completion_tokens"] == 20 diff --git a/tests/codex_lm/test_proxy_and_service_tier.py b/tests/codex_lm/test_proxy_and_service_tier.py deleted file mode 100644 index 68d7c79c..00000000 --- a/tests/codex_lm/test_proxy_and_service_tier.py +++ /dev/null @@ -1,104 +0,0 @@ -from __future__ import annotations - -import os -from unittest import mock - -import pytest -from conftest import build_stream_events -from dspy_codex_lm import CodexHTTPLM as CodexLM - - -def _patch_responses_capture(events, captured): - def fake(**kwargs): - captured["kwargs"] = kwargs - captured["https_proxy"] = os.environ.get("HTTPS_PROXY") - captured["http_proxy"] = os.environ.get("HTTP_PROXY") - return iter(events) - - return mock.patch("dspy_codex_lm.lm.litellm.responses", side_effect=fake) - - -def test_priority_service_tier_flows_into_request(): - lm = CodexLM( - model="gpt-5.5", - service_tier="priority", - access_token="fake", - account_id="fake", - ) - - request, _ = lm._build_request(prompt="x", messages=None, kwargs={}) - - assert request["service_tier"] == "priority" - - -def test_per_call_service_tier_overrides_constructor(): - lm = CodexLM( - model="gpt-5.5", - service_tier="default", - access_token="fake", - account_id="fake", - ) - - request, _ = lm._build_request( - prompt="x", messages=None, kwargs={"service_tier": "priority"} - ) - - assert request["service_tier"] == "priority" - - -@pytest.mark.parametrize( - ("env_name", "captured_name"), - [("HTTPS_PROXY", "https_proxy"), ("HTTP_PROXY", "http_proxy")], -) -def test_proxy_url_is_scoped_to_litellm_call( - monkeypatch, env_name: str, captured_name: str -): - monkeypatch.delenv("HTTPS_PROXY", raising=False) - monkeypatch.delenv("HTTP_PROXY", raising=False) - lm = CodexLM( - model="gpt-5.5", - proxy_url="http://127.0.0.1:8898", - access_token="fake", - account_id="fake", - ) - events = build_stream_events("ok", input_tokens=5, output_tokens=1) - captured: dict = {} - - with _patch_responses_capture(events, captured): - lm.forward(prompt="hi", cache=False) - - assert captured[captured_name] == "http://127.0.0.1:8898" - assert os.environ.get(env_name) is None - assert "proxy_url" not in captured["kwargs"] - - -async def test_async_proxy_url_is_scoped_to_litellm_call(monkeypatch): - monkeypatch.delenv("HTTPS_PROXY", raising=False) - monkeypatch.delenv("HTTP_PROXY", raising=False) - lm = CodexLM( - model="gpt-5.5", - proxy_url="http://127.0.0.1:8898", - access_token="fake", - account_id="fake", - ) - events = build_stream_events("ok", input_tokens=5, output_tokens=1) - captured: dict = {} - - async def _aiter(items): - for item in items: - yield item - - async def fake(**kwargs): - captured["kwargs"] = kwargs - captured["https_proxy"] = os.environ.get("HTTPS_PROXY") - captured["http_proxy"] = os.environ.get("HTTP_PROXY") - return _aiter(events) - - with mock.patch("dspy_codex_lm.lm.litellm.aresponses", side_effect=fake): - await lm.aforward(prompt="hi", cache=False) - - assert captured["https_proxy"] == "http://127.0.0.1:8898" - assert captured["http_proxy"] == "http://127.0.0.1:8898" - assert os.environ.get("HTTPS_PROXY") is None - assert os.environ.get("HTTP_PROXY") is None - assert "proxy_url" not in captured["kwargs"] diff --git a/tests/codex_lm/test_reasoning_effort.py b/tests/codex_lm/test_reasoning_effort.py deleted file mode 100644 index efaa9c94..00000000 --- a/tests/codex_lm/test_reasoning_effort.py +++ /dev/null @@ -1,151 +0,0 @@ -"""reasoning_effort is DSPy/LiteLLM's standard knob for GPT-5 / o-series -reasoning control. DSPy's request converter rewrites it to -``reasoning: {effort, summary}`` for the Responses API. These tests verify -that users can set it either on the CodexLM constructor or per-call, and -that it reaches the actual litellm.responses / aresponses call. -""" - -from unittest import mock - -import pytest -from conftest import build_stream_events -from dspy_codex_lm import CodexHTTPLM as CodexLM - -# ---- unit: _build_request rewrites effort into reasoning field ---- - - -@pytest.mark.parametrize("effort", ["low", "medium", "high"]) -def test_ctor_effort_flows_into_request(effort: str): - lm = CodexLM( - model="gpt-5.3-codex", - reasoning_effort=effort, - access_token="fake", - account_id="fake", - ) - request, _ = lm._build_request(prompt="x", messages=None, kwargs={}) - assert "reasoning_effort" not in request - assert request["reasoning"] == {"effort": effort, "summary": "auto"} - - -def test_per_call_effort_flows_into_request(lm): - request, _ = lm._build_request( - prompt="x", messages=None, kwargs={"reasoning_effort": "high"} - ) - assert "reasoning_effort" not in request - assert request["reasoning"] == {"effort": "high", "summary": "auto"} - - -def test_per_call_effort_overrides_ctor_effort(): - lm = CodexLM( - model="gpt-5.3-codex", - reasoning_effort="low", - access_token="fake", - account_id="fake", - ) - request, _ = lm._build_request( - prompt="x", messages=None, kwargs={"reasoning_effort": "high"} - ) - assert request["reasoning"] == {"effort": "high", "summary": "auto"} - - -def test_no_effort_no_reasoning_field(lm): - request, _ = lm._build_request(prompt="x", messages=None, kwargs={}) - assert "reasoning" not in request - assert "reasoning_effort" not in request - - -def test_ctor_effort_none_does_not_emit_reasoning_field(): - """Explicit None on construction must not produce ``effort: null``.""" - lm = CodexLM( - model="gpt-5.3-codex", - reasoning_effort=None, - access_token="fake", - account_id="fake", - ) - request, _ = lm._build_request(prompt="x", messages=None, kwargs={}) - assert "reasoning_effort" not in request - assert "reasoning" not in request - - -def test_per_call_effort_none_does_not_emit_reasoning_field(lm): - request, _ = lm._build_request(prompt="x", messages=None, kwargs={"reasoning_effort": None}) - assert "reasoning_effort" not in request - assert "reasoning" not in request - - -def test_per_call_effort_none_clears_ctor_effort(): - """Passing None at call time should override a ctor effort (clear it).""" - lm = CodexLM( - model="gpt-5.3-codex", - reasoning_effort="high", - access_token="fake", - account_id="fake", - ) - request, _ = lm._build_request(prompt="x", messages=None, kwargs={"reasoning_effort": None}) - assert "reasoning" not in request - - -# ---- integration: effort reaches the mocked litellm call ---- - - -def _patch_responses_capture(events, captured): - def fake(**kwargs): - captured.update(kwargs) - return iter(events) - - return mock.patch("dspy_codex_lm.lm.litellm.responses", side_effect=fake) - - -def _patch_aresponses_capture(events, captured): - async def _aiter(items): - for i in items: - yield i - - async def fake(**kwargs): - captured.update(kwargs) - return _aiter(events) - - return mock.patch("dspy_codex_lm.lm.litellm.aresponses", side_effect=fake) - - -def test_forward_sends_reasoning_to_litellm(): - lm = CodexLM( - model="gpt-5.3-codex", - reasoning_effort="high", - access_token="fake", - account_id="fake", - ) - events = build_stream_events("ok", input_tokens=5, output_tokens=1) - captured: dict = {} - with _patch_responses_capture(events, captured): - lm.forward(prompt="hi") - assert captured["reasoning"] == {"effort": "high", "summary": "auto"} - assert "reasoning_effort" not in captured - - -async def test_aforward_sends_reasoning_to_litellm(): - lm = CodexLM( - model="gpt-5.3-codex", - reasoning_effort="medium", - access_token="fake", - account_id="fake", - ) - events = build_stream_events("ok", input_tokens=5, output_tokens=1) - captured: dict = {} - with _patch_aresponses_capture(events, captured): - await lm.aforward(prompt="hi") - assert captured["reasoning"] == {"effort": "medium", "summary": "auto"} - assert "reasoning_effort" not in captured - - -def test_per_call_effort_reaches_litellm(): - lm = CodexLM( - model="gpt-5.3-codex", - access_token="fake", - account_id="fake", - ) - events = build_stream_events("ok", input_tokens=5, output_tokens=1) - captured: dict = {} - with _patch_responses_capture(events, captured): - lm.forward(prompt="hi", reasoning_effort="low") - assert captured["reasoning"] == {"effort": "low", "summary": "auto"} diff --git a/tests/codex_lm/test_retry_jitter.py b/tests/codex_lm/test_retry_jitter.py deleted file mode 100644 index 61256df3..00000000 --- a/tests/codex_lm/test_retry_jitter.py +++ /dev/null @@ -1,118 +0,0 @@ -"""Jittered exponential backoff — prevents thundering-herd retries. - -Context: - With ``wait_exponential``, every retry wakes up at exactly the same - offset after a failure (t+2s, t+6s, t+14s). When N concurrent - callers all hit a rate-limit window simultaneously (e.g. 30 eval - workers hammering Codex during peak hours), their retries land in - perfect lockstep — recreating the same concentrated load on the - very endpoint that was already struggling. Backoff is supposed to - *spread* load, not re-concentrate it. - - Switching to ``wait_random_exponential`` uniformly randomizes each - wait in [0, exponential_cap_at_attempt], so N concurrent retries - fan out across the backoff window instead of arriving in lockstep. - -RED (pre-change): ``wait_exponential(multiplier=2, max=8)`` called at - ``attempt_number=2`` always returns 4.0s — deterministic, no - variance. The "multiple calls produce distinct waits" assertion - fails. - -GREEN (post-change): ``wait_random_exponential(multiplier=2, max=8)`` - at ``attempt_number=2`` returns a fresh uniform sample from [0, 4] - on each invocation — multiple calls yield a distribution of values. -""" - -from __future__ import annotations - -from dspy_codex_lm import lm as codex_lm -from tenacity import wait_random_exponential - - -def _sample_waits(wait_fn, attempt_number: int, n: int = 50) -> list[float]: - """Invoke ``wait_fn`` ``n`` times at a fixed ``attempt_number`` and - collect the returned wait durations. Tenacity wait callables accept - a ``RetryCallState`` and return the wait in seconds; we fabricate a - minimal state to avoid booting a real retry loop. - """ - - class _MinimalState: - def __init__(self, attempt_no: int): - self.attempt_number = attempt_no - self.outcome = None - self.outcome_timestamp = None - self.start_time = 0.0 - self.retry_object = None - self.args = () - self.kwargs = {} - self.fn = None - self.idle_for = 0.0 - self.next_action = None - - state = _MinimalState(attempt_number) - return [wait_fn(state) for _ in range(n)] - - -def test_retry_wait_is_jittered(monkeypatch): - """Pulling the actual wait function from the production retry kwargs - and sampling it 50 times at attempt=2 must show variance > 0 — - deterministic ``wait_exponential`` yields a single repeated value, - jittered ``wait_random_exponential`` yields a uniform spread. - - The conftest auto-fixture zeroes ``CODEX_STREAM_WAIT_MAX`` to keep - unrelated tests fast. Restore a non-zero cap here so the jitter - has something to spread over. - """ - monkeypatch.setattr(codex_lm, "CODEX_STREAM_WAIT_MULTIPLIER", 2.0) - monkeypatch.setattr(codex_lm, "CODEX_STREAM_WAIT_MAX", 8.0) - - kwargs = codex_lm._codex_retry_kwargs() - wait_fn = kwargs["wait"] - - samples = _sample_waits(wait_fn, attempt_number=2, n=50) - - # Jitter test: at least 10 distinct values out of 50 samples. With a - # uniform [0, 4] distribution, the probability of fewer than 10 - # distinct values is effectively zero; with a deterministic function - # there's only ever one distinct value. - distinct = len(set(samples)) - assert distinct >= 10, ( - f"expected ≥10 distinct wait values across 50 samples (indicates " - f"jittered backoff), got {distinct}: {sorted(set(samples))[:5]}..." - ) - - -def test_retry_wait_respects_exponential_cap(monkeypatch): - """The jittered wait should still honor ``CODEX_STREAM_WAIT_MAX`` — - a [0, cap] uniform sample never exceeds the cap. Guards against a - hypothetical future swap to an unbounded jitter strategy. - """ - monkeypatch.setattr(codex_lm, "CODEX_STREAM_WAIT_MULTIPLIER", 2.0) - monkeypatch.setattr(codex_lm, "CODEX_STREAM_WAIT_MAX", 8.0) - - kwargs = codex_lm._codex_retry_kwargs() - wait_fn = kwargs["wait"] - - # Sample at a high attempt_number where a deterministic - # ``wait_exponential`` would have saturated at ``max`` (8.0) long - # ago. For ``wait_random_exponential`` the cap is the upper bound - # of the uniform distribution. - samples = _sample_waits(wait_fn, attempt_number=10, n=100) - - assert max(samples) <= 8.0 + 1e-6, ( - f"wait exceeded configured max of 8.0s: max observed = {max(samples)}" - ) - assert min(samples) >= 0, f"wait went negative: min = {min(samples)}" - - -def test_retry_wait_fallback_is_jitter_type(): - """Source-anchor: absent a server retry-after delay, the production - retry config must still use a jittered wait strategy. - """ - kwargs = codex_lm._codex_retry_kwargs() - wait_fn = kwargs["wait"] - assert isinstance(wait_fn._fallback, wait_random_exponential), ( - "retry fallback regressed to deterministic wait_exponential — " - "this re-enables the thundering-herd failure mode that the jitter " - "swap was introduced to prevent" - ) diff --git a/tests/codex_lm/test_stream_errors.py b/tests/codex_lm/test_stream_errors.py index bebf11f2..89f1f4dd 100644 --- a/tests/codex_lm/test_stream_errors.py +++ b/tests/codex_lm/test_stream_errors.py @@ -8,7 +8,7 @@ from unittest import mock import pytest -from conftest import make_completed, make_text_delta +from conftest import make_text_delta from dspy_codex_lm import CodexStreamError @@ -41,43 +41,6 @@ def _patch_responses(events): ) -# ---- unit: _handle_event captures failure into state ---- - - -def test_handle_event_captures_failed(lm): - state = lm._fresh_state() - lm._handle_event(_failed_event("rate_limit_exceeded", "Too many requests"), state) - assert state["failure"] is not None - assert state["failure"]["kind"] == "failed" - assert state["failure"]["code"] == "rate_limit_exceeded" - assert state["failure"]["message"] == "Too many requests" - assert state["completed"] is False - - -def test_handle_event_captures_incomplete(lm): - state = lm._fresh_state() - lm._handle_event(_incomplete_event("max_output_tokens"), state) - assert state["failure"]["kind"] == "incomplete" - assert state["failure"]["code"] == "max_output_tokens" - - -def test_handle_event_captures_error(lm): - state = lm._fresh_state() - lm._handle_event(_error_event("server_error", "internal"), state) - assert state["failure"]["kind"] == "error" - assert state["failure"]["code"] == "server_error" - - -def test_handle_event_sets_completed_flag(lm): - state = lm._fresh_state() - lm._handle_event(make_completed(input_tokens=1, output_tokens=1), state) - assert state["completed"] is True - assert state["failure"] is None - - -# ---- forward() should raise, not return empty ---- - - def test_forward_raises_on_failed_event(lm): events = [ make_text_delta("partial"), @@ -112,57 +75,3 @@ def test_forward_raises_on_truncated_stream(lm): with _patch_responses(events): with pytest.raises(CodexStreamError, match="without.*completed"): lm.forward(prompt="hi") - - -def test_forward_raises_on_empty_stream(lm): - with _patch_responses([]): - with pytest.raises(CodexStreamError): - lm.forward(prompt="hi") - - -# ---- logger emits a warning before raising ---- - - -def test_failure_logs_warning(lm, caplog): - import logging as _logging - - caplog.set_level(_logging.WARNING, logger="dspy_codex_lm.lm") - events = [_failed_event("rate_limit_exceeded", "slow down")] - with _patch_responses(events), pytest.raises(CodexStreamError): - lm.forward(prompt="hi") - records = [r for r in caplog.records if r.name == "dspy_codex_lm.lm"] - assert records - assert records[0].levelname == "WARNING" - assert "rate_limit_exceeded" in records[0].getMessage() - - -# ---- async parity ---- - - -def _async_iter(items): - async def _gen(): - for item in items: - yield item - - return _gen() - - -def _patch_aresponses(events): - async def fake(**_): - return _async_iter(events) - - return mock.patch("dspy_codex_lm.lm.litellm.aresponses", side_effect=fake) - - -async def test_aforward_raises_on_failed_event(lm): - events = [_failed_event("rate_limit_exceeded", "Too many requests")] - with _patch_aresponses(events): - with pytest.raises(CodexStreamError, match="rate_limit_exceeded"): - await lm.aforward(prompt="hi") - - -async def test_aforward_raises_on_truncated_stream(lm): - events = [make_text_delta("partial")] - with _patch_aresponses(events): - with pytest.raises(CodexStreamError, match="without.*completed"): - await lm.aforward(prompt="hi") diff --git a/tests/codex_lm/test_stream_heartbeat.py b/tests/codex_lm/test_stream_heartbeat.py index 42ce1983..f7730bb7 100644 --- a/tests/codex_lm/test_stream_heartbeat.py +++ b/tests/codex_lm/test_stream_heartbeat.py @@ -1,124 +1,47 @@ -"""RED-GREEN repro for the SSE-stream deadlock. - -Background: - ``CodexLM.aforward`` iterates the response stream via ``async for - event in stream:``. If the underlying HTTP transport stalls — - headers received, no body delivered, no FIN, no error — that loop - blocks indefinitely below Python's asyncio cancellation layer. An - outer ``asyncio.wait_for(task_timeout=...)`` cannot cancel a hung - socket read that never hits an ``await`` point, and the kernel's - TCP keepalive is 2 hours by default on macOS. - - This caused a real production stall on 2026-04-18: a SpreadsheetBench - eval seized at 10:37 with 30 concurrent workers all parked inside - ``async for event in stream:``. ``wait_for(300s)`` fired ``cancel()`` - at 10:42 but the tasks never received the CancelledError. The whole - process became a CPU-0% zombie until manually killed. - -RED (pre-fix): ``test_async_stream_that_never_yields_raises_within_heartbeat`` - runs to the outer safety-net timeout (3s) instead of raising - CodexStreamError, failing the pytest.raises assertion with a - ``asyncio.TimeoutError`` — proof the heartbeat guard is absent. - -GREEN (post-fix): each ``__anext__`` is bounded by - ``CODEX_STREAM_HEARTBEAT_SEC``. A silent stream raises - ``CodexStreamError`` within that window. - -NOTE on sync path: the sync ``forward()`` suffers the same bug but a -test for it requires subprocess isolation (a hung ``for event in stream:`` -in a pytest-owned thread kills the whole test runner). Covering the -sync path is tracked separately — the async test in this file exercises -the surface the production eval actually uses. -""" - -from __future__ import annotations - import asyncio from unittest import mock import pytest +from conftest import build_stream_events from dspy_codex_lm import CodexStreamError -from dspy_codex_lm import lm as codex_lm - -# Safety net: without the fix, the aforward call hangs forever. Wrapping -# in ``asyncio.wait_for`` gives the test a clean, bounded failure path -# (TimeoutError != CodexStreamError → the pytest.raises assertion fails -# loudly) instead of hanging the whole pytest process. -_TEST_SAFETY_TIMEOUT = 3.0 - - -class _NeverYieldsAsync: - """Async iterator that simulates a fully-hung SSE stream: ``__anext__`` - parks forever in ``asyncio.sleep``, replicating the production - deadlock where the socket is established, headers are delivered, but - no event body ever arrives and the server never closes the connection. - """ - - def __aiter__(self): - return self - - async def __anext__(self): - await asyncio.sleep(3600) # one hour — simulates "never" - raise AssertionError("unreachable") # pragma: no cover - - -async def test_default_stream_heartbeat_matches_upstream_idle_timeout(): - assert codex_lm.CODEX_STREAM_HEARTBEAT_SEC == 300.0 - -async def test_async_stream_that_never_yields_raises_within_heartbeat(lm, monkeypatch): - """The async stream consumer must raise CodexStreamError within the - configured heartbeat window when the stream goes silent. - RED state (no fix): heartbeat knob doesn't exist, the monkeypatch is - skipped, the inner ``async for event in stream:`` hangs forever, - the outer ``asyncio.wait_for(3s)`` safety net fires a TimeoutError - → the ``pytest.raises(CodexStreamError)`` assertion fails. +async def test_silent_stream_raises_instead_of_hanging(lm, monkeypatch): + monkeypatch.setattr("dspy_codex_lm.lm.CODEX_STREAM_HEARTBEAT_SEC", 0.05) - GREEN state (with fix): heartbeat=0.1s fires, tenacity retries - (disabled to 1 attempt in tests), CodexStreamError propagates - cleanly well under the safety net. - """ - # Shorten the heartbeat if the fix is present. Test stays meaningful - # in RED state (falls through to the safety net). - if hasattr(codex_lm, "CODEX_STREAM_HEARTBEAT_SEC"): - monkeypatch.setattr(codex_lm, "CODEX_STREAM_HEARTBEAT_SEC", 0.1) + async def fake_aresponses(**_): + async def stream(): + await asyncio.Future() + yield - async def _fake_aresponses(**_): - return _NeverYieldsAsync() + return stream() - with mock.patch("dspy_codex_lm.lm.litellm.aresponses", side_effect=_fake_aresponses): + with mock.patch("dspy_codex_lm.lm.litellm.aresponses", side_effect=fake_aresponses): with pytest.raises(CodexStreamError, match="stalled"): - await asyncio.wait_for( - lm.aforward(prompt="hello"), - timeout=_TEST_SAFETY_TIMEOUT, - ) + await asyncio.wait_for(lm.aforward(prompt="hi"), timeout=3) -async def test_heartbeat_does_not_interrupt_healthy_stream(lm, monkeypatch): - """Streams that emit events within the heartbeat window pass through - unchanged — the heartbeat is a ceiling, not a floor. This guards - against the fix regressing to a too-aggressive timeout that would - break normal-latency responses. +async def test_completion_does_not_wait_for_stream_disconnect(lm, monkeypatch): + monkeypatch.setattr("dspy_codex_lm.lm.CODEX_STREAM_HEARTBEAT_SEC", 0.05) - Passes in both RED and GREEN states (no-heartbeat == permissive), - so it anchors the "healthy streams still work" property. - """ - if hasattr(codex_lm, "CODEX_STREAM_HEARTBEAT_SEC"): - monkeypatch.setattr(codex_lm, "CODEX_STREAM_HEARTBEAT_SEC", 1.0) + async def fake_aresponses(**_): + async def stream(): + for event in build_stream_events("ok", input_tokens=5, output_tokens=1): + yield event + await asyncio.Future() - from conftest import build_stream_events + return stream() - events = build_stream_events("ok", input_tokens=10, output_tokens=1) + with mock.patch("dspy_codex_lm.lm.litellm.aresponses", side_effect=fake_aresponses): + response = await asyncio.wait_for(lm.aforward(prompt="hi"), timeout=3) + assert response.output[0].content[0].text == "ok" - async def _async_iter(): - for e in events: - await asyncio.sleep(0.01) # well under the 1.0s heartbeat - yield e - async def _fake_aresponses(**_): - return _async_iter() +def test_sync_completion_stops_before_reading_another_http_event(lm): + def stream(): + yield from build_stream_events("ok", input_tokens=5, output_tokens=1) + raise AssertionError("read past response.completed") - with mock.patch("dspy_codex_lm.lm.litellm.aresponses", side_effect=_fake_aresponses): - resp = await lm.aforward(prompt="hi") - assert resp.output[0].content[0].text == "ok" + with mock.patch("dspy_codex_lm.lm.litellm.responses", return_value=stream()): + response = lm.forward(prompt="hi") + assert response.output[0].content[0].text == "ok" diff --git a/tests/codex_lm/test_stream_redaction.py b/tests/codex_lm/test_stream_redaction.py new file mode 100644 index 00000000..745127e3 --- /dev/null +++ b/tests/codex_lm/test_stream_redaction.py @@ -0,0 +1,35 @@ +import json +from types import SimpleNamespace +from unittest import mock + +import pytest +from dspy_codex_lm import CodexStreamError + +from predict_rlm.debug import reset_debug_logger_for_tests + + +def test_stream_failure_logs_do_not_expose_credentials(lm, monkeypatch, tmp_path): + log_path = tmp_path / "codex-debug.jsonl" + for name in ("RLM_DEBUG", "CODEX_LM_DEBUG_LOG"): + monkeypatch.delenv(name, raising=False) + monkeypatch.setenv("PREDICT_RLM_DEBUG", "1") + monkeypatch.setenv("PREDICT_RLM_DEBUG_JSON", "1") + monkeypatch.setenv("PREDICT_RLM_DEBUG_LOG", str(log_path)) + event = SimpleNamespace( + type="response.failed", + response=SimpleNamespace( + error=SimpleNamespace(code="rate_limit_exceeded", message="slow down"), + ), + ) + reset_debug_logger_for_tests() + try: + with mock.patch("dspy_codex_lm.lm.litellm.responses", return_value=iter([event])): + with pytest.raises(CodexStreamError): + lm.forward(prompt="hi", cache=False) + text = log_path.read_text() + records = [json.loads(line) for line in text.splitlines()] + assert any(record["event"] == "codex_lm.stream.error" for record in records) + assert "fake-access" not in text + assert "fake-account" not in text + finally: + reset_debug_logger_for_tests() diff --git a/tests/codex_lm/test_stream_timing_debug.py b/tests/codex_lm/test_stream_timing_debug.py deleted file mode 100644 index 2b9a5070..00000000 --- a/tests/codex_lm/test_stream_timing_debug.py +++ /dev/null @@ -1,285 +0,0 @@ -import json -from types import SimpleNamespace -from unittest import mock - -import pytest -from conftest import make_completed, make_text_delta -from dspy_codex_lm import CodexStreamError - -from predict_rlm.debug import reset_debug_logger_for_tests - - -@pytest.fixture(autouse=True) -def reset_debug_logging(monkeypatch): - for name in ( - "PREDICT_RLM_DEBUG", - "RLM_DEBUG", - "PREDICT_RLM_DEBUG_LOG", - "PREDICT_RLM_DEBUG_JSON", - ): - monkeypatch.delenv(name, raising=False) - reset_debug_logger_for_tests() - yield - reset_debug_logger_for_tests() - - -def _patch_monotonic(monkeypatch, values): - ticks = iter(values) - monkeypatch.setattr("dspy_codex_lm.lm.monotonic", lambda: next(ticks)) - - -def _debug_records(log_path): - return [json.loads(line) for line in log_path.read_text().splitlines()] - - -def _records_by_event(log_path): - return {record["event"]: record for record in _debug_records(log_path)} - - -def _enable_json_debug(monkeypatch, tmp_path): - log_path = tmp_path / "codex-debug.jsonl" - monkeypatch.setenv("PREDICT_RLM_DEBUG", "1") - monkeypatch.setenv("PREDICT_RLM_DEBUG_JSON", "1") - monkeypatch.setenv("PREDICT_RLM_DEBUG_LOG", str(log_path)) - return log_path - - -def _response_created_event(): - return SimpleNamespace(type="response.created") - - -def test_forward_emits_stream_timing_events(lm, monkeypatch, tmp_path): - log_path = _enable_json_debug(monkeypatch, tmp_path) - _patch_monotonic(monkeypatch, [10.0, 10.1, 10.25, 10.4, 10.5, 10.55]) - events = [ - _response_created_event(), - make_text_delta("ok"), - make_completed(input_tokens=5, output_tokens=1, cached_tokens=3), - ] - - lm.kwargs["reasoning_effort"] = "xhigh" - lm.kwargs["service_tier"] = "priority" - with mock.patch("dspy_codex_lm.lm.litellm.responses", return_value=iter(events)): - result = lm.forward(prompt="hi", cache=False) - - assert result.output[0].content[0].text == "ok" - records = _records_by_event(log_path) - assert set(records) == { - "codex_lm.stream.start", - "codex_lm.stream.first_event", - "codex_lm.stream.first_text_delta", - "codex_lm.stream.end", - } - assert records["codex_lm.stream.start"]["attempt_number"] == 1 - assert records["codex_lm.stream.start"]["model"] == "gpt-5.3-codex" - assert records["codex_lm.stream.start"]["transport"] == "http_sse" - assert records["codex_lm.stream.start"]["reasoning_effort"] == "xhigh" - assert records["codex_lm.stream.start"]["service_tier"] == "priority" - assert records["codex_lm.stream.first_event"]["reasoning_effort"] == "xhigh" - assert records["codex_lm.stream.first_text_delta"]["service_tier"] == "priority" - assert records["codex_lm.stream.end"]["reasoning_effort"] == "xhigh" - assert records["codex_lm.stream.end"]["service_tier"] == "priority" - assert records["codex_lm.stream.first_event"]["first_event_type"] == "response.created" - assert records["codex_lm.stream.first_event"]["tt_first_event_ms"] == 100.0 - assert records["codex_lm.stream.first_text_delta"]["ttft_ms"] == 250.0 - assert records["codex_lm.stream.end"]["stream_total_ms"] == 500.0 - assert records["codex_lm.stream.end"]["parse_overhead_ms"] == 50.0 - assert records["codex_lm.stream.end"]["output_text_chars"] == 2 - assert records["codex_lm.stream.end"]["completed"] is True - assert records["codex_lm.stream.end"]["prompt_tokens"] == 5 - assert records["codex_lm.stream.end"]["cached_prompt_tokens"] == 3 - assert records["codex_lm.stream.end"]["prompt_cache_read_ratio"] == pytest.approx(0.6) - - -async def test_aforward_emits_stream_timing_events(lm, monkeypatch, tmp_path): - log_path = _enable_json_debug(monkeypatch, tmp_path) - _patch_monotonic(monkeypatch, [20.0, 20.2, 20.4, 20.6, 20.7, 20.72]) - events = [ - _response_created_event(), - make_text_delta("async"), - make_completed(input_tokens=5, output_tokens=1), - ] - - async def fake_aresponses(**_): - async def gen(): - for event in events: - yield event - - return gen() - - with mock.patch("dspy_codex_lm.lm.litellm.aresponses", side_effect=fake_aresponses): - result = await lm.aforward(prompt="hi", cache=False) - - assert result.output[0].content[0].text == "async" - records = _records_by_event(log_path) - assert records["codex_lm.stream.first_event"]["tt_first_event_ms"] == 200.0 - assert records["codex_lm.stream.first_text_delta"]["ttft_ms"] == 400.0 - assert records["codex_lm.stream.end"]["stream_total_ms"] == 700.0 - assert records["codex_lm.stream.end"]["output_text_chars"] == 5 - assert records["codex_lm.stream.end"]["completed"] is True - - -def test_codex_debug_log_gets_records_when_predict_rlm_debug_is_also_enabled( - lm, - monkeypatch, - tmp_path, -): - predict_log = tmp_path / "predict-debug.jsonl" - codex_log = tmp_path / "codex-debug.jsonl" - monkeypatch.setenv("PREDICT_RLM_DEBUG", "1") - monkeypatch.setenv("PREDICT_RLM_DEBUG_JSON", "1") - monkeypatch.setenv("PREDICT_RLM_DEBUG_LOG", str(predict_log)) - monkeypatch.setenv("CODEX_LM_DEBUG_LOG", str(codex_log)) - _patch_monotonic(monkeypatch, [10.0, 10.1, 10.2, 10.3, 10.35, 10.36]) - events = [make_text_delta("ok"), make_completed(input_tokens=5, output_tokens=1)] - - with mock.patch("dspy_codex_lm.lm.litellm.responses", return_value=iter(events)): - lm.forward(prompt="hi", cache=False) - - records = _records_by_event(codex_log) - assert records["codex_lm.stream.start"]["transport"] == "http_sse" - assert records["codex_lm.stream.end"]["transport"] == "http_sse" - assert not predict_log.exists() - - - -def test_stream_end_reports_non_text_event_counts_and_inter_event_gap(lm, monkeypatch, tmp_path): - log_path = _enable_json_debug(monkeypatch, tmp_path) - _patch_monotonic( - monkeypatch, - [ - 40.0, # stream start - 40.1, # response.created - 40.4, # response.in_progress - 41.0, # first text delta - 41.2, # response.completed - 41.25, # mark stream end - 41.27, # emit end - ], - ) - events = [ - _response_created_event(), - SimpleNamespace(type="response.in_progress"), - make_text_delta("ok"), - make_completed(input_tokens=5, output_tokens=1), - ] - - with mock.patch("dspy_codex_lm.lm.litellm.responses", return_value=iter(events)): - lm.forward(prompt="hi", cache=False) - - end = _records_by_event(log_path)["codex_lm.stream.end"] - assert end["events_before_first_text"] == 2 - assert end["max_inter_event_gap_ms"] == 600.0 - assert end["non_text_event_counts"] == { - "response.created": 1, - "response.in_progress": 1, - "response.completed": 1, - } - - -def test_stream_events_include_selected_rotation_profile(monkeypatch, tmp_path): - import json - from pathlib import Path - - from dspy_codex_lm import CodexHTTPLM as CodexLM - from dspy_codex_lm.auth import import_auth_profile - from dspy_codex_lm.cli import main - - def write_auth(path: Path, *, access_token: str, account_id: str) -> Path: - path.write_text( - json.dumps( - { - "tokens": { - "access_token": access_token, - "account_id": account_id, - } - } - ), - encoding="utf-8", - ) - return path - - log_path = _enable_json_debug(monkeypatch, tmp_path) - monkeypatch.setattr(Path, "home", classmethod(lambda cls: tmp_path)) - import_auth_profile( - "alpha", - write_auth(tmp_path / "alpha.json", access_token="alpha-token", account_id="acct-alpha"), - ) - import_auth_profile( - "beta", - write_auth(tmp_path / "beta.json", access_token="beta-token", account_id="acct-beta"), - ) - assert main(["codex-lm", "rotation", "on"]) == 0 - - def choose_credentials(credentials): - return next( - credential - for credential in credentials - if credential.account_id == "acct-beta" - ) - - monkeypatch.setattr("dspy_codex_lm.lm.random.choice", choose_credentials) - _patch_monotonic(monkeypatch, [60.0, 60.1, 60.2, 60.3, 60.35, 60.36]) - events = [make_text_delta("ok"), make_completed(input_tokens=5, output_tokens=1)] - - with mock.patch("dspy_codex_lm.lm.litellm.responses", return_value=iter(events)): - CodexLM(model="gpt-5.3-codex").forward(prompt="hi", cache=False) - - records = _records_by_event(log_path) - assert records["codex_lm.stream.start"]["auth_profile"] == "beta" - assert records["codex_lm.stream.start"]["auth_source"] == "rotation" - assert records["codex_lm.stream.end"]["auth_profile"] == "beta" - - -def test_stream_error_reports_last_event_when_stream_stalls(monkeypatch, tmp_path): - log_path = _enable_json_debug(monkeypatch, tmp_path) - from dspy_codex_lm.lm import _StreamTiming - - timing = _StreamTiming( - model="gpt-5.5", - attempt_number=1, - start_at=50.0, - transport="websocket", - ) - _patch_monotonic(monkeypatch, [50.2, 80.2]) - timing.observe_event(_response_created_event()) - timing.emit_error({}, CodexStreamError("stall")) - - error = _records_by_event(log_path)["codex_lm.stream.error"] - assert error["transport"] == "websocket" - assert error["last_event_type"] == "response.created" - assert error["last_event_age_ms"] == 30000.0 - assert error["events_before_first_text"] == 1 - assert error["max_inter_event_gap_ms"] is None - - -def test_stream_error_event_is_sanitized(lm, monkeypatch, tmp_path): - log_path = _enable_json_debug(monkeypatch, tmp_path) - _patch_monotonic(monkeypatch, [30.0, 30.1, 30.2, 30.3]) - events = [ - make_text_delta("partial"), - SimpleNamespace( - type="response.failed", - response=SimpleNamespace( - error=SimpleNamespace( - code="rate_limit_exceeded", - message="slow down", - ), - ), - ), - ] - - with mock.patch("dspy_codex_lm.lm.litellm.responses", return_value=iter(events)): - with pytest.raises(CodexStreamError): - lm.forward(prompt="hi", cache=False) - - records = _records_by_event(log_path) - error = records["codex_lm.stream.error"] - assert error["failure_kind"] == "failed" - assert error["failure_code"] == "rate_limit_exceeded" - assert error["exception_type"] == "CodexStreamError" - assert error["completed"] is False - assert error["output_text_chars"] == len("partial") - log_text = log_path.read_text() - assert "fake-access" not in log_text - assert "fake-account" not in log_text diff --git a/tests/codex_lm/test_terminal_completion.py b/tests/codex_lm/test_terminal_completion.py deleted file mode 100644 index 2a712999..00000000 --- a/tests/codex_lm/test_terminal_completion.py +++ /dev/null @@ -1,65 +0,0 @@ -from __future__ import annotations - -import asyncio -import threading -from unittest import mock - -from conftest import build_stream_events - - -class _CompletedThenPendingSync: - def __init__(self, events): - self._events = iter(events) - self.pending_reads = 0 - - def __iter__(self): - return self - - def __next__(self): - try: - return next(self._events) - except StopIteration: - self.pending_reads += 1 - threading.Event().wait(0.25) - raise StopIteration - - -class _CompletedThenPendingAsync: - def __init__(self, events): - self._events = iter(events) - self.pending_reads = 0 - - def __aiter__(self): - return self - - async def __anext__(self): - try: - return next(self._events) - except StopIteration: - self.pending_reads += 1 - await asyncio.Future() - raise AssertionError("unreachable") # pragma: no cover - - -def test_forward_stops_at_response_completed_before_pending_stream(lm, monkeypatch): - monkeypatch.setattr("dspy_codex_lm.lm.CODEX_STREAM_HEARTBEAT_SEC", 0.05) - stream = _CompletedThenPendingSync(build_stream_events("ok", input_tokens=5, output_tokens=1)) - - with mock.patch("dspy_codex_lm.lm.litellm.responses", return_value=stream): - response = lm.forward(prompt="hi") - - assert response.output[0].content[0].text == "ok" - assert stream.pending_reads == 0 - - -async def test_aforward_stops_at_response_completed_before_pending_stream(lm, monkeypatch): - monkeypatch.setattr("dspy_codex_lm.lm.CODEX_STREAM_HEARTBEAT_SEC", 0.05) - stream = _CompletedThenPendingAsync( - build_stream_events("ok", input_tokens=5, output_tokens=1) - ) - - with mock.patch("dspy_codex_lm.lm.litellm.aresponses", return_value=stream): - response = await asyncio.wait_for(lm.aforward(prompt="hi"), timeout=1.0) - - assert response.output[0].content[0].text == "ok" - assert stream.pending_reads == 0 diff --git a/tests/codex_lm/test_usage.py b/tests/codex_lm/test_usage.py index ec683582..25739fe4 100644 --- a/tests/codex_lm/test_usage.py +++ b/tests/codex_lm/test_usage.py @@ -1,391 +1,53 @@ -import json -from pathlib import Path +from dspy_codex_lm.usage import format_profile_usage_summaries, summarize_usage -from dspy_codex_lm.usage import ( - CODEX_USAGE_ENDPOINT, - fetch_codex_usage, - format_disabled_profile_usage_entry, - format_profile_usage_summaries, - format_usage_summary, - summarize_usage, -) - -def _auth_file(tmp_path: Path) -> Path: - path = tmp_path / "auth.json" - path.write_text( - json.dumps( - { - "tokens": { - "access_token": "secret-token", - "account_id": "acct-secret", - "refresh_token": "secret-refresh", - }, - "user": {"email": "person@example.com"}, - } - ) - ) - return path - - -def test_fetch_codex_usage_builds_required_headers(tmp_path: Path): - seen = {} - - def fake_transport(url, *, headers, timeout): - seen["url"] = url - seen["headers"] = headers - seen["timeout"] = timeout - return {"rate_limit": {}} - - payload = fetch_codex_usage( - auth_path=_auth_file(tmp_path), - transport=fake_transport, - ) - - assert payload == {"rate_limit": {}} - assert seen == { - "url": CODEX_USAGE_ENDPOINT, - "headers": { - "Authorization": "Bearer secret-token", - "ChatGPT-Account-ID": "acct-secret", - "Accept": "application/json", - }, - "timeout": 10.0, - } - - -def test_summarize_usage_parses_rate_limit_credit_and_model_windows(): - payload = { - "rate_limit": { - "primary": { - "used": 12, - "limit": 40, - "remaining": 28, - "reset_at": "2026-05-10T12:00:00Z", - }, - "secondary": { - "current_value": 9, - "max_value": 10, - "resets_in_seconds": 1800, - }, - }, - "credits": { - "granted": 1000, - "used": 250, - "expires_at": "2026-06-01T00:00:00Z", - }, - "additional": { - "models": { - "gpt-5.3-codex": { - "rate_limit": { - "remaining": 3, - "limit": 20, - "reset_after_seconds": 600, - } - } - } - }, - "user": {"email": "person@example.com"}, - "account_id": "acct-secret", - } - - rows = summarize_usage(payload) - - assert [row.label for row in rows] == [ - "additional.models.gpt-5.3-codex.rate_limit", - "credits", - "rate_limit.primary", - "rate_limit.secondary", - ] - assert rows[0].percent_remaining == 15.0 - assert rows[0].reset == "in 10m" - assert rows[1].remaining == 750 - assert rows[1].percent_remaining == 75.0 - assert rows[1].reset == "2026-06-01T00:00:00Z" - assert rows[2].remaining == 28 - assert rows[2].percent_remaining == 70.0 - assert rows[2].reset == "2026-05-10T12:00:00Z" - assert rows[3].remaining == 1 - assert rows[3].percent_remaining == 10.0 - assert rows[3].reset == "in 30m" - - -def test_summarize_usage_parses_live_wham_shape(): +def test_usage_derives_remaining_credit_and_nested_model_limits(): payload = { - "plan_type": "pro", - "rate_limit": { - "allowed": True, - "limit_reached": False, - "primary_window": { - "used_percent": 8, - "limit_window_seconds": 18000, - "reset_after_seconds": 5719, - "reset_at": 1778387445, - }, - "secondary_window": { - "used_percent": 35, - "limit_window_seconds": 604800, - "reset_after_seconds": 157475, - "reset_at": 1778539200, - }, - }, - "additional_rate_limits": [ - { - "limit_name": "GPT-5.3-Codex-Spark", - "metered_feature": "codex.spark", - "rate_limit": { - "primary_window": { - "used_percent": 0, - "limit_window_seconds": 18000, - "reset_after_seconds": 18000, - }, - "secondary_window": { - "used_percent": 0, - "limit_window_seconds": 604800, - "reset_after_seconds": 604800, - }, - }, - } - ], - "credits": { - "balance": "0", - "has_credits": False, - "unlimited": False, - "approx_cloud_messages": [0, 0], - }, - "user_id": "user-secret", + "credits": {"granted": 1000, "used": 250}, + "rate_limit": {"primary": {"current_value": 9, "max_value": 10}}, + "models": {"codex": {"rate_limit": {"remaining": 3, "limit": 20}}}, } + rows = {row.label: row for row in summarize_usage(payload)} + assert rows["credits"].remaining == 750 + assert rows["credits"].percent_remaining == 75 + assert rows["rate_limit.primary"].remaining == 1 + assert rows["rate_limit.primary"].percent_remaining == 10 + assert rows["models.codex.rate_limit"].percent_remaining == 15 - rows = summarize_usage(payload) - assert [row.label for row in rows] == [ - "GPT-5.3-Codex-Spark.primary_window", - "GPT-5.3-Codex-Spark.secondary_window", - "rate_limit.primary_window", - "rate_limit.secondary_window", - ] - assert rows[0].percent_remaining == 100.0 - assert rows[0].reset == "in 5h" - assert rows[1].percent_remaining == 100.0 - assert rows[1].reset == "in 7d" - assert rows[2].percent_remaining == 92.0 - assert rows[2].reset == "04:30 on 10 May" - assert rows[3].percent_remaining == 65.0 - assert rows[3].reset == "22:40 on 11 May" - - text = format_usage_summary(payload) - assert "Plan: pro" in text - assert "Credits: balance 0; has_credits=false" in text - assert ( - "5h limit: [██████████████████░░] 92% left (resets 04:30 on 10 May)" - ) in text - assert ( - "Weekly limit: [█████████████░░░░░░░] 65% left (resets 22:40 on 11 May)" - ) in text - assert "GPT-5.3-Codex-Spark limit:" in text - assert ( - " 5h limit: [████████████████████] 100% left (resets in 5h)" - ) in text - assert ( - " Weekly limit: [████████████████████] 100% left (resets in 7d)" - ) in text - - -def test_format_usage_summary_groups_top_level_live_windows_under_general_header(): +def test_live_usage_windows_are_not_conflated_with_model_specific_limits(): payload = { "rate_limit": { - "primary_window": { - "used_percent": 12, - "limit_window_seconds": 18000, - "reset_after_seconds": 300, - }, - "secondary_window": { - "used_percent": 35, - "limit_window_seconds": 604800, - "reset_after_seconds": 604800, - }, - }, - "additional_rate_limits": [ - { - "limit_name": "GPT-5.3-Codex-Spark", - "rate_limit": { - "primary_window": { - "used_percent": 0, - "limit_window_seconds": 18000, - "reset_after_seconds": 18000, - }, - }, - } - ], - } - - text = format_usage_summary(payload) - - lines = text.splitlines() - assert lines[0] == "-" * 60 - assert lines[-1] == "-" * 60 - assert lines[1:-1] == [ - "Codex usage", - "General usage limits:", - " 5h limit: [██████████████████░░] 88% left (resets in 5m)", - " Weekly limit: [█████████████░░░░░░░] 65% left (resets in 7d)", - "GPT-5.3-Codex-Spark limit:", - " 5h limit: [████████████████████] 100% left (resets in 5h)", - ] - - -def test_format_profile_usage_summaries_groups_general_and_model_limits_per_profile(): - work_payload = { - "rate_limit": { - "primary_window": { - "used_percent": 12, - "limit_window_seconds": 18000, - "reset_after_seconds": 300, - }, + "primary_window": {"used_percent": 8, "limit_window_seconds": 18000}, + "secondary_window": {"used_percent": 35, "limit_window_seconds": 604800}, }, "additional_rate_limits": [ { "limit_name": "GPT-5.3-Codex-Spark", "rate_limit": { - "primary_window": { - "used_percent": 0, - "limit_window_seconds": 18000, - "reset_after_seconds": 18000, - }, + "primary_window": {"used_percent": 0, "limit_window_seconds": 18000}, }, } ], } - personal_payload = { - "rate_limit": { - "secondary_window": { - "used_percent": 35, - "limit_window_seconds": 604800, - "reset_after_seconds": 604800, - }, - }, - "additional_rate_limits": [ - { - "limit_name": "GPT-5.4", - "rate_limit": { - "primary_window": { - "used_percent": 20, - "limit_window_seconds": 18000, - "reset_after_seconds": 60, - }, - }, - } - ], + rows = {row.label: row.percent_remaining for row in summarize_usage(payload)} + assert rows == { + "rate_limit.primary_window": 92, + "rate_limit.secondary_window": 65, + "GPT-5.3-Codex-Spark.primary_window": 100, } - text = format_profile_usage_summaries( - [("work", work_payload), ("personal", personal_payload)], - default_profile="personal", - ) - - lines = text.splitlines() - assert lines[0] == "-" * 60 - assert lines[-1] == "-" * 60 - assert lines[1:-1] == [ - "work:", - " General usage limits:", - " 5h limit: [██████████████████░░] 88% left (resets in 5m)", - " GPT-5.3-Codex-Spark limit:", - " 5h limit: [████████████████████] 100% left (resets in 5h)", - "", - "personal (default):", - " General usage limits:", - " Weekly limit: [█████████████░░░░░░░] 65% left (resets in 7d)", - " GPT-5.4 limit:", - " 5h limit: [████████████████░░░░] 80% left (resets in 1m)", - ] - -def test_format_profile_usage_summaries_preserves_profile_email_display_name(): +def test_profile_summary_preserves_display_name_without_leaking_payload_secrets(): payload = { - "rate_limit": { - "primary_window": { - "used_percent": 12, - "limit_window_seconds": 18000, - "reset_after_seconds": 300, - }, - }, + "rate_limit": {"primary": {"used": 1, "limit": 4}}, "user": {"email": "payload@example.com"}, - } - - text = format_profile_usage_summaries([("gabriel@example.com", payload)]) - colored = format_profile_usage_summaries( - [("gabriel@example.com", payload)], - color=True, - ) - - assert "gabriel@example.com:" in text - assert "[redacted-email]:" not in text - assert "payload@example.com" not in text - assert "\x1b[1;36mgabriel@example.com:\x1b[0m" in colored - assert "[redacted-email]:" not in colored - assert "payload@example.com" not in colored - - -def test_format_disabled_profile_usage_entry_is_labeled_and_redacted(): - text = format_profile_usage_summaries( - [ - ("work", {"rate_limit": {"primary": {"used": 1, "limit": 2}}}), - format_disabled_profile_usage_entry("acct-secret"), - ], - default_profile="acct-secret", - ) - - assert "work:" in text - assert "[redacted-account] (disabled) (default):" in text - assert " Disabled; live usage fetch skipped." in text - assert "acct-secret" not in text - - -def test_format_usage_summary_color_can_be_enabled_and_disabled(): - payload = { - "rate_limit": { - "primary_window": { - "used_percent": 12, - "limit_window_seconds": 18000, - "reset_after_seconds": 300, - }, - }, - } - - plain = format_usage_summary(payload, color=False) - colored = format_usage_summary(payload, color=True) - - assert "\x1b[" not in plain - assert "\x1b[" in colored - assert "General usage limits:" in colored - assert "[██████████████████░░] 88% left" in colored - - -def test_format_usage_summary_is_redacted_and_stable(): - payload = { - "rate_limit": { - "primary": { - "used": 1, - "limit": 4, - "remaining": 3, - "reset_at": "2026-05-10T12:00:00Z", - } - }, "account_id": "acct-secret", - "email": "person@example.com", "access_token": "secret-token", } - - text = format_usage_summary(payload) - - lines = text.splitlines() - assert lines[0] == "-" * 60 - assert lines[-1] == "-" * 60 - assert lines[1:-1] == [ - "Codex usage", - "rate_limit.primary: 3/4 remaining (75.0% remaining); resets 2026-05-10T12:00:00Z", - ] - assert "secret-token" not in text + text = format_profile_usage_summaries([("profile@example.com", payload)]) + assert "profile@example.com" in text + assert "75.0%" in text + assert "payload@example.com" not in text assert "acct-secret" not in text - assert "person@example.com" not in text + assert "secret-token" not in text diff --git a/tests/codex_lm/test_usage_tracker_hook.py b/tests/codex_lm/test_usage_tracker_hook.py deleted file mode 100644 index 003f5616..00000000 --- a/tests/codex_lm/test_usage_tracker_hook.py +++ /dev/null @@ -1,166 +0,0 @@ -"""CodexLM fires ``dspy.settings.usage_tracker.add_usage`` after each call. - -DSPy's LM base class fires this at ``clients/lm.py:167`` (sync) and :205 -(async) so that ``dspy.track_usage()`` can attribute tokens to the -prediction. CodexLM overrides ``forward``/``aforward`` entirely, which -before this fix silently bypassed the hook. Downstream consumers -(predict-rlm's ``pred.get_lm_usage()`` accumulation, cost_log, etc.) -then saw $0 for every CodexLM-routed call. - -These tests pin the hook contract: - - Within a ``track_usage()`` context, each CodexLM call populates the - tracker under the LM's ``self.model`` slug. - - Cache hits DO NOT fire the hook (matching DSPy's behavior). - - Without a tracker in context, the call succeeds silently. - -Real tests exercise the full CodexLM.aforward path with litellm mocked -at the responses layer — not the LM method itself — so the hook's -guard conditions run against real response objects, not MagicMocks -that would silently make ``dict(usage)`` empty. -""" - -import copy -from unittest import mock - -import dspy -from conftest import build_stream_events -from dspy.utils.usage_tracker import UsageTracker - - -def test_aforward_populates_usage_tracker(lm): - events = build_stream_events("hi", input_tokens=1000, output_tokens=50) - - import asyncio - - async def _fake_aresponses(**_): - async def _gen(): - for ev in copy.deepcopy(events): - yield ev - - return _gen() - - tracker = UsageTracker() - - async def run(): - with dspy.settings.context(usage_tracker=tracker): - with mock.patch( - "dspy_codex_lm.lm.litellm.aresponses", - side_effect=_fake_aresponses, - ): - await lm.aforward(prompt="hi") - - asyncio.run(run()) - totals = tracker.get_total_tokens() - assert lm.model in totals, ( - f"expected tracker to have {lm.model!r}, got {list(totals.keys())}" - ) - assert totals[lm.model]["prompt_tokens"] == 1000 - assert totals[lm.model]["completion_tokens"] == 50 - - -def test_forward_populates_usage_tracker(lm): - events = build_stream_events("hi", input_tokens=1000, output_tokens=50) - - def _fake_responses(**_): - return iter(copy.deepcopy(events)) - - tracker = UsageTracker() - - with dspy.settings.context(usage_tracker=tracker): - with mock.patch( - "dspy_codex_lm.lm.litellm.responses", - side_effect=_fake_responses, - ): - lm.forward(prompt="hi") - - totals = tracker.get_total_tokens() - assert totals[lm.model]["prompt_tokens"] == 1000 - assert totals[lm.model]["completion_tokens"] == 50 - - -def test_no_tracker_in_context_call_still_succeeds(lm): - """When no ``usage_tracker`` is set in context, the hook does nothing - and the call proceeds normally — no crash, no exception. - """ - events = build_stream_events("hi", input_tokens=1000, output_tokens=50) - - def _fake_responses(**_): - return iter(copy.deepcopy(events)) - - # No usage_tracker set - with mock.patch( - "dspy_codex_lm.lm.litellm.responses", - side_effect=_fake_responses, - ): - resp = lm.forward(prompt="hi") - assert resp.output[0].content[0].text == "hi" - - -def test_track_usage_via_context_populates_pred_lm_usage(lm): - """The end-to-end integration: wrap ``dspy.Predict`` in - ``track_usage=True``, predict once, confirm ``pred.get_lm_usage()`` - shows the CodexLM's tokens. This is the exact path predict-rlm's - cost accounting depends on. - """ - import asyncio - - class Sig(dspy.Signature): - q: str = dspy.InputField() - a: str = dspy.OutputField() - - dspy.settings.configure(lm=lm) - predict = dspy.Predict(Sig) - - # Completion text in the content must be a valid JSON/Chat response - # the adapter can parse — we use JSON-format output for simplicity. - events = build_stream_events('{"a": "42"}', input_tokens=500, output_tokens=20) - - async def _fake_aresponses(**_): - async def _gen(): - for ev in copy.deepcopy(events): - yield ev - - return _gen() - - async def run(): - with dspy.settings.context(track_usage=True): - with mock.patch( - "dspy_codex_lm.lm.litellm.aresponses", - side_effect=_fake_aresponses, - ): - pred = await predict.acall(q="hi") - return pred.get_lm_usage() - - usage = asyncio.run(run()) - assert usage, f"pred.get_lm_usage() should be populated, got {usage!r}" - assert lm.model in usage - assert usage[lm.model]["prompt_tokens"] > 0 - - -def test_cache_hit_skips_usage_tracker_hook(lm, monkeypatch): - """A cached response MUST NOT populate the tracker — no new tokens - were consumed. Mirrors DSPy's own ``not getattr(results, "cache_hit", - False)`` guard. - """ - events = build_stream_events("hi", input_tokens=1000, output_tokens=50) - - def _fake_responses(**_): - return iter(copy.deepcopy(events)) - - tracker = UsageTracker() - - with dspy.settings.context(usage_tracker=tracker): - with mock.patch( - "dspy_codex_lm.lm.litellm.responses", - side_effect=_fake_responses, - ): - # Warm the cache - lm.forward(prompt="same prompt") - # Second call = cache hit - lm.forward(prompt="same prompt") - - totals = tracker.get_total_tokens() - # Only the first (non-cache-hit) call populated tokens — not 2× - assert totals[lm.model]["prompt_tokens"] == 1000, ( - f"cache hit should not double-count: got {totals}" - ) diff --git a/tests/codex_lm/test_ws_lm.py b/tests/codex_lm/test_ws_lm.py index e99643ea..6ae8e284 100644 --- a/tests/codex_lm/test_ws_lm.py +++ b/tests/codex_lm/test_ws_lm.py @@ -1,220 +1,76 @@ -from __future__ import annotations - import copy -from dataclasses import dataclass from types import SimpleNamespace from unittest import mock -import pytest from conftest import build_stream_events from dspy_codex_lm import CodexHTTPLM, CodexStreamError, CodexWSLM -@dataclass -class RecordedTurn: - request_id: str - sticky_state: dict - request: dict - headers: dict - - class FakeWSTransport: def __init__(self, streams): - self.streams = list(streams) - self.turns: list[RecordedTurn] = [] + self.streams = iter(streams) + self.turns = [] def stream_turn(self, *, request, headers, request_id, sticky_state): - self.turns.append( - RecordedTurn( - request_id=request_id, - sticky_state=sticky_state, - request=copy.deepcopy(request), - headers=copy.deepcopy(headers), - ) - ) - stream = self.streams.pop(0) + self.turns.append((request_id, sticky_state)) + stream = next(self.streams) if isinstance(stream, BaseException): raise stream return iter(copy.deepcopy(stream)) - async def astream_turn(self, *, request, headers, request_id, sticky_state): - self.turns.append( - RecordedTurn( - request_id=request_id, - sticky_state=sticky_state, - request=copy.deepcopy(request), - headers=copy.deepcopy(headers), - ) - ) - stream = self.streams.pop(0) - if isinstance(stream, BaseException): - raise stream - - async def _events(): - for event in copy.deepcopy(stream): - yield event - - return _events() - - -class FakeHTTPFallback: - def __init__(self): - self.forward_calls = [] - self.aforward_calls = [] - - def forward(self, prompt=None, messages=None, **kwargs): - self.forward_calls.append((prompt, messages, dict(kwargs))) - return SimpleNamespace(source="http", prompt=prompt) - - async def aforward(self, prompt=None, messages=None, **kwargs): - self.aforward_calls.append((prompt, messages, dict(kwargs))) - return SimpleNamespace(source="http", prompt=prompt) - - -def _failed_events(code: str = "rate_limit_exceeded", msg: str = "slow down"): - return [ - SimpleNamespace( - type="response.failed", - response=SimpleNamespace( - error=SimpleNamespace(code=code, message=msg), - ), - ) - ] - - -def test_codex_ws_lm_is_exported_and_accepts_codex_lm_kwargs(): - lm = CodexWSLM( - model="gpt-5.3-codex", - instructions="You are Ada.", - access_token="fake-access", - account_id="fake-account", - proxy_url="http://127.0.0.1:9999", - reasoning_effort="low", - ) - - request, _ = lm._build_request(prompt="hi", messages=None, kwargs={}) - assert request["instructions"] == "You are Ada." - assert request["reasoning"]["effort"] == "low" - - -def test_two_forward_calls_create_distinct_ws_turns(): - transport = FakeWSTransport( - [ - build_stream_events("one", response_id="resp-one"), - build_stream_events("two", response_id="resp-two"), - ] - ) - lm = CodexWSLM( - model="gpt-5.3-codex", - access_token="fake-access", - account_id="fake-account", - ws_transport=transport, - ) - - first = lm.forward(prompt="one", cache=False) - second = lm.forward(prompt="two", cache=False) - assert first.output[0].content[0].text == "one" - assert second.output[0].content[0].text == "two" - assert len(transport.turns) == 2 - assert transport.turns[0].request_id != transport.turns[1].request_id - assert transport.turns[0].headers["session_id"] != transport.turns[1].headers["session_id"] - assert transport.turns[0].sticky_state is not transport.turns[1].sticky_state - - -def test_retry_stays_inside_one_invocation_turn_boundary(monkeypatch): +def test_retry_shares_turn_state_without_leaking_it_to_next_request(monkeypatch): monkeypatch.setattr("dspy_codex_lm.lm.CODEX_STREAM_MAX_ATTEMPTS", 2) + failed = SimpleNamespace( + type="response.failed", + response=SimpleNamespace( + error=SimpleNamespace(code="rate_limit_exceeded", message="slow down") + ), + ) transport = FakeWSTransport( [ - _failed_events(), - build_stream_events("ok", response_id="resp-ok"), - build_stream_events("next", response_id="resp-next"), + [failed], + build_stream_events("recovered"), + build_stream_events("next"), ] ) lm = CodexWSLM( model="gpt-5.3-codex", - access_token="fake-access", - account_id="fake-account", + access_token="fake", + account_id="fake", ws_transport=transport, + ws_fallback=False, ) + assert lm.forward(prompt="retry", cache=False).output[0].content[0].text == "recovered" + assert lm.forward(prompt="next", cache=False).output[0].content[0].text == "next" + first, retry, subsequent = transport.turns + assert first[0] == retry[0] + assert first[1] is retry[1] + assert subsequent[0] != retry[0] + assert subsequent[1] is not retry[1] - first = lm.forward(prompt="retry", cache=False) - second = lm.forward(prompt="next", cache=False) - assert first.output[0].content[0].text == "ok" - assert second.output[0].content[0].text == "next" - assert len(transport.turns) == 3 - assert transport.turns[0].request_id == transport.turns[1].request_id - assert transport.turns[0].sticky_state is transport.turns[1].sticky_state - assert transport.turns[1].request_id != transport.turns[2].request_id - assert transport.turns[1].sticky_state is not transport.turns[2].sticky_state - - -def test_ws_fallback_exhaustion_routes_later_invocations_to_http(monkeypatch): - monkeypatch.setattr("dspy_codex_lm.lm.CODEX_STREAM_MAX_ATTEMPTS", 1) +def test_exhausted_websocket_stays_on_http_for_later_invocations(): transport = FakeWSTransport([CodexStreamError("ws unavailable")]) - fallback = FakeHTTPFallback() + fallback = CodexHTTPLM(model="gpt-5.3-codex", access_token="fake", account_id="fake") lm = CodexWSLM( model="gpt-5.3-codex", - access_token="fake-access", - account_id="fake-account", + access_token="fake", + account_id="fake", ws_transport=transport, fallback_lm=fallback, ) - - first = lm.forward(prompt="first", cache=False) - second = lm.forward(prompt="second", cache=False) - - assert first.source == "http" - assert second.source == "http" - assert [call[0] for call in fallback.forward_calls] == ["first", "second"] + with mock.patch( + "dspy_codex_lm.lm.litellm.responses", + side_effect=[ + iter(build_stream_events("first", input_tokens=5, output_tokens=1)), + iter(build_stream_events("second", input_tokens=9, output_tokens=2)), + ], + ): + first = lm.forward(prompt="first", cache=False) + second = lm.forward(prompt="second", cache=False) + assert first.output[0].content[0].text == "first" + assert first.usage.input_tokens == 5 + assert second.output[0].content[0].text == "second" + assert second.usage.input_tokens == 9 assert len(transport.turns) == 1 - - -async def test_two_aforward_calls_create_distinct_ws_turns(): - transport = FakeWSTransport( - [ - build_stream_events("one", response_id="resp-one"), - build_stream_events("two", response_id="resp-two"), - ] - ) - lm = CodexWSLM( - model="gpt-5.3-codex", - access_token="fake-access", - account_id="fake-account", - ws_transport=transport, - ) - - first = await lm.aforward(prompt="one", cache=False) - second = await lm.aforward(prompt="two", cache=False) - - assert first.output[0].content[0].text == "one" - assert second.output[0].content[0].text == "two" - assert len(transport.turns) == 2 - assert transport.turns[0].request_id != transport.turns[1].request_id - assert transport.turns[0].sticky_state is not transport.turns[1].sticky_state - - -@pytest.mark.parametrize("lm_kind", ["http", "ws"]) -def test_protocol_forward_response_contract_is_shared(lm_kind): - events = build_stream_events("contract", input_tokens=100, output_tokens=1) - if lm_kind == "http": - lm = CodexHTTPLM( - model="gpt-5.3-codex", - access_token="fake-access", - account_id="fake-account", - ) - with mock.patch("dspy_codex_lm.lm.litellm.responses", return_value=iter(events)): - resp = lm.forward(prompt="shared", cache=False) - else: - lm = CodexWSLM( - model="gpt-5.3-codex", - access_token="fake-access", - account_id="fake-account", - ws_transport=FakeWSTransport([events]), - ) - resp = lm.forward(prompt="shared", cache=False) - - assert resp.output[0].content[0].text == "contract" - assert resp.usage.input_tokens == 100 - assert resp.usage.cost > 0 diff --git a/tests/fixtures/bootstrap_controller/debian13-no-python/Dockerfile b/tests/fixtures/bootstrap_controller/debian13-no-python/Dockerfile deleted file mode 100644 index 1e93cd1d..00000000 --- a/tests/fixtures/bootstrap_controller/debian13-no-python/Dockerfile +++ /dev/null @@ -1,21 +0,0 @@ -FROM debian:13.0-slim - -ENV ROOT=/tmp/predict_rlm_controller - -RUN apt-get update \ - && DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends ca-certificates \ - && rm -rf /var/lib/apt/lists/* - -COPY src/predict_rlm/remote/bootstrap_controller.sh /tmp/bootstrap_controller.sh - -RUN mkdir -p "$ROOT/repo" \ - && printf '%s\n' 'def ok(): return "ok"' > "$ROOT/repo/fake_controller_repo.py" \ - && printf '%s\n' \ - 'from setuptools import setup' \ - 'setup(name="fake-controller-repo", version="0.0.0", py_modules=["fake_controller_repo"], extras_require={"codex-lm": []})' \ - > "$ROOT/repo/setup.py" - -RUN sh /tmp/bootstrap_controller.sh --root "$ROOT" --repo "$ROOT/repo" --extra '[codex-lm]' --python 3.12 > /tmp/bootstrap.log 2>&1 \ - || { cat /tmp/bootstrap.log; exit 1; } \ - && grep -q 'repairing Python/pip/venv with package manager: apt' /tmp/bootstrap.log \ - && "$ROOT/.venv/bin/python" -c 'import sys, fake_controller_repo; assert sys.version_info[:2] == (3, 12), sys.version; assert fake_controller_repo.ok() == "ok"' diff --git a/tests/fixtures/bootstrap_controller/python311-slim/Dockerfile b/tests/fixtures/bootstrap_controller/python311-slim/Dockerfile deleted file mode 100644 index 853a3d3e..00000000 --- a/tests/fixtures/bootstrap_controller/python311-slim/Dockerfile +++ /dev/null @@ -1,17 +0,0 @@ -FROM python:3.11-slim - -ENV ROOT=/tmp/predict_rlm_controller - -COPY src/predict_rlm/remote/bootstrap_controller.sh /tmp/bootstrap_controller.sh - -RUN mkdir -p "$ROOT/repo" \ - && printf '%s\n' 'def ok(): return "ok"' > "$ROOT/repo/fake_controller_repo.py" \ - && printf '%s\n' \ - 'from setuptools import setup' \ - 'setup(name="fake-controller-repo", version="0.0.0", py_modules=["fake_controller_repo"], extras_require={"codex-lm": []})' \ - > "$ROOT/repo/setup.py" - -RUN sh /tmp/bootstrap_controller.sh --root "$ROOT" --repo "$ROOT/repo" --extra '[codex-lm]' --python 3.12 > /tmp/bootstrap.log 2>&1 \ - || { cat /tmp/bootstrap.log; exit 1; } \ - && ! grep -q 'repairing Python/pip/venv' /tmp/bootstrap.log \ - && "$ROOT/.venv/bin/python" -c 'import sys, fake_controller_repo; assert sys.version_info[:2] == (3, 12), sys.version; assert fake_controller_repo.ok() == "ok"' diff --git a/tests/fixtures/bootstrap_controller/tbench-python-313/Dockerfile b/tests/fixtures/bootstrap_controller/tbench-python-313/Dockerfile deleted file mode 100644 index bbd5acf4..00000000 --- a/tests/fixtures/bootstrap_controller/tbench-python-313/Dockerfile +++ /dev/null @@ -1,16 +0,0 @@ -FROM ghcr.io/laude-institute/t-bench/python-3-13:20250620 - -ENV ROOT=/tmp/predict_rlm_controller - -COPY src/predict_rlm/remote/bootstrap_controller.sh /tmp/bootstrap_controller.sh - -RUN mkdir -p "$ROOT/repo" \ - && printf '%s\n' 'def ok(): return "ok"' > "$ROOT/repo/fake_controller_repo.py" \ - && printf '%s\n' \ - 'from setuptools import setup' \ - 'setup(name="fake-controller-repo", version="0.0.0", py_modules=["fake_controller_repo"], extras_require={"codex-lm": []})' \ - > "$ROOT/repo/setup.py" - -RUN sh /tmp/bootstrap_controller.sh --root "$ROOT" --repo "$ROOT/repo" --extra '[codex-lm]' --python 3.12 > /tmp/bootstrap.log 2>&1 \ - || { cat /tmp/bootstrap.log; exit 1; } \ - && "$ROOT/.venv/bin/python" -c 'import sys, fake_controller_repo; assert sys.version_info[:2] == (3, 12), sys.version; assert fake_controller_repo.ok() == "ok"' diff --git a/tests/fixtures/bootstrap_controller/ubuntu24-python-pip-no-venv/Dockerfile b/tests/fixtures/bootstrap_controller/ubuntu24-python-pip-no-venv/Dockerfile deleted file mode 100644 index 2c748d4f..00000000 --- a/tests/fixtures/bootstrap_controller/ubuntu24-python-pip-no-venv/Dockerfile +++ /dev/null @@ -1,21 +0,0 @@ -FROM ubuntu:24.04 - -ENV ROOT=/tmp/predict_rlm_controller - -RUN apt-get update \ - && DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends ca-certificates python3 python3-pip \ - && rm -rf /var/lib/apt/lists/* - -COPY src/predict_rlm/remote/bootstrap_controller.sh /tmp/bootstrap_controller.sh - -RUN mkdir -p "$ROOT/repo" \ - && printf '%s\n' 'def ok(): return "ok"' > "$ROOT/repo/fake_controller_repo.py" \ - && printf '%s\n' \ - 'from setuptools import setup' \ - 'setup(name="fake-controller-repo", version="0.0.0", py_modules=["fake_controller_repo"], extras_require={"codex-lm": []})' \ - > "$ROOT/repo/setup.py" - -RUN sh /tmp/bootstrap_controller.sh --root "$ROOT" --repo "$ROOT/repo" --extra '[codex-lm]' --python 3.12 > /tmp/bootstrap.log 2>&1 \ - || { cat /tmp/bootstrap.log; exit 1; } \ - && grep -q 'repairing Python/pip/venv with package manager: apt' /tmp/bootstrap.log \ - && "$ROOT/.venv/bin/python" -c 'import sys, fake_controller_repo; assert sys.version_info[:2] == (3, 12), sys.version; assert fake_controller_repo.ok() == "ok"' diff --git a/tests/runtime_contracts/backends.py b/tests/runtime_contracts/backends.py index a57de427..64dc75ee 100644 --- a/tests/runtime_contracts/backends.py +++ b/tests/runtime_contracts/backends.py @@ -1,222 +1,86 @@ from __future__ import annotations -import importlib import os import shutil +import socket import sys -from dataclasses import dataclass, field +from dataclasses import dataclass from pathlib import Path -from typing import Any, Callable, Literal, Protocol +from typing import Any, Callable import pytest -# The runtime-contract matrix exercises the supervisor/SBX backend (local-supervisor -# seam and real SBX), so the whole package requires the [sbx] extra (websockets). -pytest.importorskip("websockets") +from predict_rlm.backends import DirectPythonBackend, JspiBackend -from predict_rlm.backends import ( # noqa: E402 - DirectPythonBackend, - JspiBackend, - SbxBackend, - SbxConfig, +PAYLOAD_PATH = ( + Path(__file__).resolve().parents[2] / "src/predict_rlm/backends/supervisor/_payload.py" ) -ROOT = Path(__file__).resolve().parents[2] -TERMINAL_BENCH_DIR = ROOT / "examples" / "terminal_bench" -if str(TERMINAL_BENCH_DIR) not in sys.path: - sys.path.insert(0, str(TERMINAL_BENCH_DIR)) - -runner_module = importlib.import_module("terminal_bench_rlm.tools.runner") -runner_script_path = runner_module.runner_script_path - - -CAPABILITIES = frozenset( - { - "execute", - "state", - "reset", - "code_fences", - "submit", - "deferred_submit", - "recoverable_errors", - "host_tools", - "recoverable_iteration_timeout", - "files", - } -) - - -class RuntimeHandle(Protocol): - spec: RuntimeSpec - - def require(self, capability: str) -> None: ... - - def configure( - self, - *, - tools: dict[str, Callable[..., Any]] | None = None, - output_fields: list[dict[str, Any]] | None = None, - ) -> None: ... - - def execute(self, code: str, *, timeout: float | None = None) -> Any: ... - def output(self, result: Any) -> str: ... - - def timeout_observation(self, result: Any) -> dict[str, Any]: ... - - def defer_next_submit_finalization(self) -> None: ... +def _predict_tool(signature: str, **kwargs: Any) -> dict[str, Any]: + return {"answer": "4"} - def reset(self) -> None: ... - def shutdown(self) -> None: ... +def _shape_tool(kind: str) -> Any: + return {"list": [1, 2], "dict": {"ok": True}, "none": None, "text": "hello"}[kind] - def mount_file_at(self, host_path: str, sandbox_path: str) -> None: ... - def mkdir_p(self, sandbox_path: str) -> None: ... +def _failing_tool() -> None: + raise ValueError("host tool failed") - def list_dir(self, sandbox_path: str) -> list[str]: ... - def sync_file_to(self, sandbox_path: str, host_path: str) -> None: ... +def _default_tools() -> dict[str, Callable[..., Any]]: + return {"predict": _predict_tool, "shape_tool": _shape_tool, "failing_tool": _failing_tool} @dataclass(frozen=True) class RuntimeSpec: name: str - adapter: Literal[ - "jspi-process", - "sbx-cli", - "direct-process", - "test-only-local-supervisor", - ] - environment: Literal[ - "deno-subprocess", - "sbx-sandbox", - "direct-process", - "local-supervisor-seam", - ] - engine: Literal["pyodide-jspi", "python-runner"] - make: Callable[[Path, "RuntimeSpec"], RuntimeHandle] - capabilities: frozenset[str] - opt_in: bool = False - skip_reason: str | None = None - xfail_contracts: dict[str, str] = field(default_factory=dict) - - -def _predict_tool(signature: str, **kwargs: Any) -> dict[str, Any]: - del signature, kwargs - return {"answer": "4"} - - -def _shape_tool(kind: str) -> Any: - if kind == "list": - return [1, 2] - if kind == "dict": - return {"ok": True} - if kind == "none": - return None - if kind == "text": - return "hello" - raise ValueError(f"unknown shape: {kind}") - - -def _failing_tool() -> str: - raise ValueError("host tool failed") - + make: Callable[[Path, "RuntimeSpec"], "RuntimeHandle"] + unsupported: frozenset[str] = frozenset() -def _default_tools() -> dict[str, Callable[..., Any]]: - return { - "predict": _predict_tool, - "shape_tool": _shape_tool, - "failing_tool": _failing_tool, - } +class RuntimeHandle: + """Normalize the two legacy result/reset interfaces, not backend behavior.""" -class InterpreterRuntimeHandle: def __init__(self, spec: RuntimeSpec, interpreter: Any) -> None: self.spec = spec self.interpreter = interpreter + def __getattr__(self, name: str) -> Any: + return getattr(self.interpreter, name) + def require(self, capability: str) -> None: - if capability not in CAPABILITIES: - raise AssertionError(f"unknown runtime capability: {capability}") - if capability in self.spec.xfail_contracts: - pytest.xfail(self.spec.xfail_contracts[capability]) - if capability not in self.spec.capabilities: - pytest.skip( - self.spec.skip_reason - or f"{self.spec.name} does not advertise {capability}" - ) - - def configure( - self, - *, - tools: dict[str, Callable[..., Any]] | None = None, - output_fields: list[dict[str, Any]] | None = None, - ) -> None: - configure_runtime = getattr(self.interpreter, "configure_runtime", None) - if configure_runtime is None: - pytest.skip(f"{self.spec.name} does not support runtime reconfiguration") - configure_runtime(tools=tools, output_fields=output_fields) - - def execute(self, code: str, *, timeout: float | None = None) -> Any: - return self.interpreter.execute(code, timeout=timeout) - - def output(self, result: Any) -> str: - if isinstance(result, str): - return result - if isinstance(result, dict) and "output" in result: - return str(result["output"]) - return str(result) - - def timeout_observation(self, result: Any) -> dict[str, Any]: - if isinstance(result, dict) and "timeout" in result: - return { - "seconds": result["timeout"]["seconds"], - "stdout": result.get("stdout", ""), - "stderr": result.get("stderr", ""), - "state": result.get("state"), - } - return { - "seconds": getattr(result, "timeout_seconds"), - "stdout": getattr(result, "stdout", ""), - "stderr": getattr(result, "stderr", ""), - "state": getattr(result, "state", None), - } + if capability in self.spec.unsupported: + pytest.skip(f"{self.spec.name} does not support {capability}") - def defer_next_submit_finalization(self) -> None: - defer = getattr(self.interpreter, "defer_next_submit_finalization", None) - if defer is None: - pytest.skip(f"{self.spec.name} does not support deferred submit") - defer() + def configure(self, **kwargs: Any) -> None: + self.interpreter.configure_runtime(**kwargs) def reset(self) -> None: - reset = getattr(self.interpreter, "reset", None) - if reset is None: + if isinstance(self.interpreter, JspiBackend): self.interpreter.shutdown() - return - reset() - - def shutdown(self) -> None: - self.interpreter.shutdown() - - def mount_file_at(self, host_path: str, sandbox_path: str) -> None: - self.interpreter.mount_file_at(host_path, sandbox_path) + else: + self.interpreter.reset() - def mkdir_p(self, sandbox_path: str) -> None: - self.interpreter.mkdir_p(sandbox_path) + @staticmethod + def output(result: Any) -> str: + return result["output"] if isinstance(result, dict) else str(result) - def list_dir(self, sandbox_path: str) -> list[str]: - return self.interpreter.list_dir(sandbox_path) - - def sync_file_to(self, sandbox_path: str, host_path: str) -> None: - self.interpreter.sync_file_to(sandbox_path, host_path) + @staticmethod + def timeout_observation(result: Any) -> dict[str, Any]: + return { + "seconds": result.timeout_seconds, + "stdout": result.stdout, + "stderr": result.stderr, + "state": result.state, + } def _make_jspi(tmp_path: Path, spec: RuntimeSpec) -> RuntimeHandle: - del tmp_path if shutil.which("deno") is None: pytest.skip("JSPI contracts require Deno") - return InterpreterRuntimeHandle( + return RuntimeHandle( spec, JspiBackend( tools=_default_tools(), @@ -226,8 +90,8 @@ def _make_jspi(tmp_path: Path, spec: RuntimeSpec) -> RuntimeHandle: ) -def _make_direct_process(tmp_path: Path, spec: RuntimeSpec) -> RuntimeHandle: - return InterpreterRuntimeHandle( +def _make_direct(tmp_path: Path, spec: RuntimeSpec) -> RuntimeHandle: + return RuntimeHandle( spec, DirectPythonBackend( tools=_default_tools(), @@ -240,86 +104,55 @@ def _make_direct_process(tmp_path: Path, spec: RuntimeSpec) -> RuntimeHandle: def _make_sbx(tmp_path: Path, spec: RuntimeSpec) -> RuntimeHandle: - if os.environ.get("PREDICT_RLM_RUN_SBX_TESTS") != "1": + if spec.name == "sbx" and ( + os.environ.get("PREDICT_RLM_RUN_SBX_TESTS") != "1" or shutil.which("sbx") is None + ): pytest.skip( - "real SBX runtime contracts require PREDICT_RLM_RUN_SBX_TESTS=1, " - "the sbx CLI, and sbx login" + "real SBX contracts require PREDICT_RLM_RUN_SBX_TESTS=1, sbx CLI, and login" ) - if shutil.which("sbx") is None: - pytest.skip("real SBX runtime contracts require the sbx CLI") - return InterpreterRuntimeHandle( + pytest.importorskip("websockets") + from predict_rlm.backends import SbxBackend, SbxConfig + + kwargs: dict[str, Any] = {} + if spec.name == "sbx/local-websocket": + with socket.socket() as sock: + sock.bind(("127.0.0.1", 0)) + port = sock.getsockname()[1] + path = f"/runtime-contract-{os.getpid()}-{port}" + kwargs = { + "_websocket_supervisor_command": [ + sys.executable, + "-u", + str(PAYLOAD_PATH), + "--websocket-host", + "127.0.0.1", + "--websocket-port", + str(port), + "--websocket-path", + path, + ], + "_websocket_url": f"ws://127.0.0.1:{port}{path}", + } + return RuntimeHandle( spec, SbxBackend( config=SbxConfig(name="runtime-contract-sbx", exec_timeout=10), tools=_default_tools(), preinstall_packages=False, _staging_root=tmp_path / "sbx-staging", - ), - ) - - -def _make_internal_jsonrpc(tmp_path: Path, spec: RuntimeSpec) -> RuntimeHandle: - return InterpreterRuntimeHandle( - spec, - SbxBackend( - config=SbxConfig(name="runtime-contract-local-supervisor", exec_timeout=10), - tools=_default_tools(), - preinstall_packages=False, - _supervisor_command=[sys.executable, "-u", str(runner_script_path())], - _staging_root=tmp_path / "internal-jsonrpc-staging", + **kwargs, ), ) def runtime_specs() -> list[RuntimeSpec]: return [ + RuntimeSpec("jspi", _make_jspi, frozenset({"deferred_submit"})), RuntimeSpec( - name="jspi", - adapter="jspi-process", - environment="deno-subprocess", - engine="pyodide-jspi", - make=_make_jspi, - capabilities=frozenset( - { - "execute", - "state", - "reset", - "code_fences", - "recoverable_errors", - "host_tools", - "recoverable_iteration_timeout", - } - ), - ), - RuntimeSpec( - name="python-runner/direct-process", - adapter="direct-process", - environment="direct-process", - engine="python-runner", - make=_make_direct_process, - capabilities=frozenset(CAPABILITIES), - ), - RuntimeSpec( - name="sbx", - adapter="sbx-cli", - environment="sbx-sandbox", - engine="python-runner", - make=_make_sbx, - capabilities=frozenset(CAPABILITIES), - opt_in=True, - xfail_contracts={ - "tool_timeout": ( - "real SBX per-host-tool timeout is not implemented in the " - "shared contract matrix yet" - ) - }, - ), - RuntimeSpec( - name="internal/python-runner-jsonrpc", - adapter="test-only-local-supervisor", - environment="local-supervisor-seam", - engine="python-runner", - make=_make_internal_jsonrpc, - capabilities=frozenset(CAPABILITIES), + "python-runner/direct-process", + _make_direct, + frozenset({"partial_error_output", "concurrent_tools"}), ), + RuntimeSpec("sbx/local-websocket", _make_sbx, frozenset({"deferred_submit"})), + RuntimeSpec("sbx", _make_sbx, frozenset({"deferred_submit"})), ] diff --git a/tests/runtime_contracts/conftest.py b/tests/runtime_contracts/conftest.py index c04addc0..28930560 100644 --- a/tests/runtime_contracts/conftest.py +++ b/tests/runtime_contracts/conftest.py @@ -4,25 +4,24 @@ import pytest -from .backends import RuntimeHandle, RuntimeSpec, runtime_specs +from .backends import RuntimeHandle, runtime_specs -_HERE = Path(__file__).resolve().parent +def _runtime_params(): + for spec in runtime_specs(): + marks = [] + if spec.name == "jspi": + marks.append(pytest.mark.integration) + elif spec.name.startswith("sbx"): + marks.append(pytest.mark.sbx) + if spec.name == "sbx": + marks.append(pytest.mark.integration) + yield pytest.param(spec, id=spec.name, marks=marks) -def pytest_collection_modifyitems(items: list[pytest.Item]) -> None: - """Mark every runtime-contract test as `sbx` so it runs only in the sbx CI job.""" - for item in items: - try: - if _HERE in Path(str(item.fspath)).resolve().parents: - item.add_marker(pytest.mark.sbx) - except (OSError, ValueError): - continue - -@pytest.fixture(params=runtime_specs(), ids=lambda spec: spec.name) +@pytest.fixture(params=list(_runtime_params())) def runtime(request: pytest.FixtureRequest, tmp_path: Path) -> RuntimeHandle: - spec: RuntimeSpec = request.param - handle = spec.make(tmp_path, spec) + handle = request.param.make(tmp_path, request.param) try: yield handle finally: diff --git a/tests/runtime_contracts/test_backend_matrix.py b/tests/runtime_contracts/test_backend_matrix.py deleted file mode 100644 index 00b1db66..00000000 --- a/tests/runtime_contracts/test_backend_matrix.py +++ /dev/null @@ -1,34 +0,0 @@ -from __future__ import annotations - -from .backends import CAPABILITIES, runtime_specs - - -def test_runtime_specs_have_unique_maintained_target_names() -> None: - specs = runtime_specs() - - assert [spec.name for spec in specs] == [ - "jspi", - "python-runner/direct-process", - "sbx", - "internal/python-runner-jsonrpc", - ] - assert len({spec.name for spec in specs}) == len(specs) - - -def test_runtime_specs_advertise_only_known_capabilities() -> None: - unknown = { - capability - for spec in runtime_specs() - for capability in spec.capabilities - if capability not in CAPABILITIES - } - - assert unknown == set() - - -def test_legacy_environment_api_paths_are_not_primitive_targets() -> None: - target_names = {spec.name for spec in runtime_specs()} - - assert "python-runner/environment-api" not in target_names - assert "harbor/environment-api" not in target_names - assert "docker-container-adapter" not in target_names diff --git a/tests/runtime_contracts/test_error_contract.py b/tests/runtime_contracts/test_error_contract.py deleted file mode 100644 index 7c990153..00000000 --- a/tests/runtime_contracts/test_error_contract.py +++ /dev/null @@ -1,25 +0,0 @@ -from __future__ import annotations - -import pytest -from dspy.primitives.code_interpreter import CodeInterpreterError - -from .backends import RuntimeHandle - - -def test_recoverable_user_exception_allows_later_execute(runtime: RuntimeHandle) -> None: - runtime.require("recoverable_errors") - - with pytest.raises((CodeInterpreterError, NameError)) as exc_info: - runtime.execute("raise ValueError('ordinary failure')") - - assert "ordinary failure" in str(exc_info.value) - assert runtime.output(runtime.execute("print('recovered')")) == "recovered\n" - - -def test_syntax_error_allows_later_execute(runtime: RuntimeHandle) -> None: - runtime.require("recoverable_errors") - - with pytest.raises(SyntaxError): - runtime.execute("for") - - assert runtime.output(runtime.execute("print('after syntax')")) == "after syntax\n" diff --git a/tests/runtime_contracts/test_execution_contract.py b/tests/runtime_contracts/test_execution_contract.py index af5c881c..f2dee39f 100644 --- a/tests/runtime_contracts/test_execution_contract.py +++ b/tests/runtime_contracts/test_execution_contract.py @@ -1,35 +1,122 @@ from __future__ import annotations +import asyncio +import threading +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path + import pytest +from dspy.primitives.code_interpreter import CodeInterpreterError, FinalOutput from .backends import RuntimeHandle def test_execute_preserves_state_until_reset(runtime: RuntimeHandle) -> None: - runtime.require("execute") - runtime.require("state") - runtime.require("reset") - assert runtime.output(runtime.execute("counter = 40\nprint('ready')")) == "ready\n" assert runtime.output(runtime.execute("counter += 2\nprint(counter)")) == "42\n" - runtime.reset() + assert runtime.output(runtime.execute("print('counter' in globals())")) == "False\n" + + +def test_concurrent_execute_requests_are_serialized(runtime: RuntimeHandle) -> None: + runtime.execute("counter = 0") + started = threading.Barrier(2) + + def increment(): + loop = asyncio.new_event_loop() + asyncio.set_event_loop(loop) + try: + started.wait(timeout=5) + return runtime.output( + runtime.execute( + "import time\n" + "previous = counter\n" + "time.sleep(0.05)\n" + "counter = previous + 1\n" + "print(counter)" + ) + ) + finally: + loop.close() + asyncio.set_event_loop(None) + + with ThreadPoolExecutor(max_workers=2) as executor: + first = executor.submit(increment) + second = executor.submit(increment) + assert sorted([first.result(timeout=20), second.result(timeout=20)]) == ["1\n", "2\n"] + + +@pytest.mark.parametrize("fence", ["python", "repl", ""]) +def test_code_fence_normalization(runtime: RuntimeHandle, fence: str) -> None: + assert runtime.output(runtime.execute(f"```{fence}\nprint(42)\n```")) == "42\n" + + +def test_user_exception_reports_error_and_allows_recovery( + runtime: RuntimeHandle, +) -> None: + with pytest.raises(CodeInterpreterError) as raised: + runtime.execute("print('before failure')\nraise ValueError('ordinary failure')") + if "partial_error_output" not in runtime.spec.unsupported: + assert raised.value.partial_output == "before failure\n" + assert "ordinary failure" in str(raised.value) + assert runtime.output(runtime.execute("print('recovered')")) == "recovered\n" + + +def test_syntax_error_allows_later_execute(runtime: RuntimeHandle) -> None: + with pytest.raises(SyntaxError): + runtime.execute("for") + assert runtime.output(runtime.execute("print('after syntax')")) == "after syntax\n" + + +def test_submit_returns_final_output(runtime: RuntimeHandle) -> None: + runtime.configure(output_fields=[{"name": "answer", "annotation": "str", "type": "str"}]) + result = runtime.execute("SUBMIT(answer='done')") + assert isinstance(result, FinalOutput) + assert result.output == {"answer": "done"} + + +def test_deferred_submit_keeps_runtime_alive_until_confirmed(runtime: RuntimeHandle) -> None: + runtime.require("deferred_submit") + runtime.configure(output_fields=[{"name": "answer", "annotation": "str"}]) + runtime.defer_next_submit_finalization() + deferred = runtime.execute("SUBMIT(answer='draft')") + probe = runtime.execute("print('alive after deferred submit')") + final = runtime.execute("SUBMIT(answer='confirmed')") + assert isinstance(deferred, FinalOutput) + assert deferred.output == {"answer": "draft"} + assert runtime.output(probe) == "alive after deferred submit\n" + assert isinstance(final, FinalOutput) + assert final.output == {"answer": "confirmed"} - with pytest.raises((NameError, SyntaxError, RuntimeError, Exception)) as exc_info: - runtime.execute("print(counter)") - assert "counter" in str(exc_info.value) - assert runtime.output(runtime.execute("print('after reset')")) == "after reset\n" +def test_file_operations_round_trip(runtime: RuntimeHandle, tmp_path: Path) -> None: + source = tmp_path / "input.txt" + target = tmp_path / "output.txt" + source.write_text("hello", encoding="utf-8") + runtime.mount_file_at(str(source), "/sandbox/input.txt") + runtime.mkdir_p("/sandbox/out") + runtime.execute( + "text = open('/sandbox/input.txt').read()\n" + "open('/sandbox/out/result.txt', 'w').write(text + ' world')" + ) + assert "/sandbox/out/result.txt" in runtime.list_dir("/sandbox/out") + runtime.sync_file_to("/sandbox/out/result.txt", str(target)) + assert target.read_text(encoding="utf-8") == "hello world" -@pytest.mark.parametrize( - ("source", "expected"), - [ - ("```python\nprint('python')\n```", "python\n"), - ("```repl\nprint('repl')\n```", "repl\n"), - ("```\nprint('bare')\n```", "bare\n"), - ], -) -def test_code_fence_normalization(runtime: RuntimeHandle, source: str, expected: str) -> None: - runtime.require("code_fences") - assert runtime.output(runtime.execute(source)) == expected +def test_recoverable_timeout_preserves_output_and_recovers(runtime: RuntimeHandle) -> None: + result = runtime.execute( + "import sys, time\n" + "print('before timeout')\n" + "print('stderr before timeout', file=sys.stderr)\n" + "sys.stdout.flush(); sys.stderr.flush()\n" + "while True:\n" + " time.sleep(0.05)\n", + timeout=0.2, + ) + followup = runtime.execute("print('after timeout')") + timeout = runtime.timeout_observation(result) + assert timeout["seconds"] == 0.2 + assert timeout["stdout"] == "before timeout\n" + assert timeout["stderr"].startswith("stderr before timeout\n") + assert runtime.output(followup) == "after timeout\n" diff --git a/tests/runtime_contracts/test_file_contract.py b/tests/runtime_contracts/test_file_contract.py deleted file mode 100644 index 3efdd533..00000000 --- a/tests/runtime_contracts/test_file_contract.py +++ /dev/null @@ -1,31 +0,0 @@ -from __future__ import annotations - -from pathlib import Path - -from .backends import RuntimeHandle - - -def test_basic_file_operations_round_trip( - runtime: RuntimeHandle, - tmp_path: Path, -) -> None: - runtime.require("files") - - source = tmp_path / "input.txt" - target = tmp_path / "output.txt" - source.write_text("hello", encoding="utf-8") - - runtime.mount_file_at(str(source), "/sandbox/input.txt") - runtime.mkdir_p("/sandbox/out") - result = runtime.execute( - "text = open('/sandbox/input.txt').read()\n" - "open('/sandbox/out/result.txt', 'w').write(text + ' world')\n" - "print(text)" - ) - - assert runtime.output(result) == "hello\n" - assert "/sandbox/out/result.txt" in runtime.list_dir("/sandbox/out") - - runtime.sync_file_to("/sandbox/out/result.txt", str(target)) - - assert target.read_text(encoding="utf-8") == "hello world" diff --git a/tests/runtime_contracts/test_submit_contract.py b/tests/runtime_contracts/test_submit_contract.py deleted file mode 100644 index 07a8a19d..00000000 --- a/tests/runtime_contracts/test_submit_contract.py +++ /dev/null @@ -1,31 +0,0 @@ -from __future__ import annotations - -from dspy.primitives.code_interpreter import FinalOutput - -from .backends import RuntimeHandle - - -def test_submit_returns_final_output(runtime: RuntimeHandle) -> None: - runtime.require("submit") - runtime.configure(output_fields=[{"name": "answer", "annotation": "str"}]) - - result = runtime.execute("SUBMIT(answer='done')") - - assert isinstance(result, FinalOutput) - assert result.output == {"answer": "done"} - - -def test_deferred_submit_keeps_runtime_alive_until_confirmed(runtime: RuntimeHandle) -> None: - runtime.require("deferred_submit") - runtime.configure(output_fields=[{"name": "answer", "annotation": "str"}]) - - runtime.defer_next_submit_finalization() - deferred = runtime.execute("SUBMIT(answer='draft')") - probe = runtime.execute("print('alive after deferred submit')") - final = runtime.execute("SUBMIT(answer='confirmed')") - - assert isinstance(deferred, FinalOutput) - assert deferred.output == {"answer": "draft"} - assert runtime.output(probe) == "alive after deferred submit\n" - assert isinstance(final, FinalOutput) - assert final.output == {"answer": "confirmed"} diff --git a/tests/runtime_contracts/test_timeout_contract.py b/tests/runtime_contracts/test_timeout_contract.py deleted file mode 100644 index 09601de2..00000000 --- a/tests/runtime_contracts/test_timeout_contract.py +++ /dev/null @@ -1,26 +0,0 @@ -from __future__ import annotations - -from .backends import RuntimeHandle - - -def test_recoverable_iteration_timeout_preserves_output_and_recovers( - runtime: RuntimeHandle, -) -> None: - runtime.require("recoverable_iteration_timeout") - - result = runtime.execute( - "import sys, time\n" - "print('before timeout')\n" - "print('stderr before timeout', file=sys.stderr)\n" - "sys.stdout.flush(); sys.stderr.flush()\n" - "while True:\n" - " time.sleep(0.05)\n", - timeout=0.2, - ) - followup = runtime.execute("print('after timeout')") - - timeout = runtime.timeout_observation(result) - assert timeout["seconds"] == 0.2 - assert timeout["stdout"] == "before timeout\n" - assert timeout["stderr"].startswith("stderr before timeout\n") - assert runtime.output(followup) == "after timeout\n" diff --git a/tests/runtime_contracts/test_tool_contract.py b/tests/runtime_contracts/test_tool_contract.py index d8e0081d..72c2552f 100644 --- a/tests/runtime_contracts/test_tool_contract.py +++ b/tests/runtime_contracts/test_tool_contract.py @@ -1,10 +1,12 @@ from __future__ import annotations +import asyncio import multiprocessing import os import queue as queue_module import shutil import subprocess +import threading import time from pathlib import Path @@ -12,7 +14,69 @@ from .backends import RuntimeHandle -_REAL_SBX_RUNTIME_SANDBOX = "runtime-contract-sbx" + +def test_host_tool_result_shapes(runtime: RuntimeHandle) -> None: + result = runtime.execute( + "items = await shape_tool('list')\n" + "mapping = await shape_tool('dict')\n" + "none_value = await shape_tool('none')\n" + "text = await shape_tool('text')\n" + "print(items, mapping['ok'], none_value is None, text)" + ) + assert runtime.output(result) == "[1, 2] True True hello\n" + + +@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync-tool", "async-tool"]) +def test_host_tools_run_concurrently(runtime: RuntimeHandle, asynchronous: bool) -> None: + runtime.require("concurrent_tools") + barrier = threading.Barrier(2) + arrivals = 0 + both_started = None + + def sync_tool(value): + barrier.wait(timeout=2) + return value * 2 + + async def async_tool(value): + nonlocal arrivals, both_started + if both_started is None: + both_started = asyncio.Event() + arrivals += 1 + if arrivals == 2: + both_started.set() + await asyncio.wait_for(both_started.wait(), timeout=2) + return value * 2 + + runtime.configure(tools={"double": async_tool if asynchronous else sync_tool}) + result = runtime.execute( + "import asyncio\nprint(await asyncio.gather(double(3), double(7)))" + ) + assert runtime.output(result) == "[6, 14]\n" + + +# Large pipe messages have intermittently stalled on loaded CI; retain the +# boundary regression as local rather than multiplying the same request. +@pytest.mark.local +def test_large_host_tool_request_round_trips(runtime: RuntimeHandle) -> None: + def inspect_payload(payload): + assert payload == "x" * 950_000 + return len(payload) + + runtime.configure(tools={"inspect_payload": inspect_payload}) + result = runtime.execute("print(await inspect_payload('x' * 950000))") + assert runtime.output(result) == "950000\n" + + +def test_tool_exception_allows_later_tool_use(runtime: RuntimeHandle) -> None: + result = runtime.execute( + "try:\n" + " await failing_tool()\n" + "except Exception as exc:\n" + " print('host tool failed' in str(exc))" + ) + assert runtime.output(result) == "True\n" + followup = runtime.execute("print((await predict('question -> answer'))['answer'])") + assert runtime.output(followup) == "4\n" def _slow_tool() -> str: @@ -20,209 +84,77 @@ def _slow_tool() -> str: return "slow" -def _run_concurrent_host_tool_timeout_repro( - runtime_name: str, - staging: str, - queue: multiprocessing.Queue, -) -> None: +def _run_timeout_repro(runtime_name: str, staging: str, result_queue) -> None: from .backends import runtime_specs spec = next(spec for spec in runtime_specs() if spec.name == runtime_name) runtime = None try: runtime = spec.make(Path(staging), spec) - runtime.require("host_tools") - runtime.require("recoverable_iteration_timeout") runtime.configure(tools={"slow_tool": _slow_tool}) runtime.execute("pass") - queue.put(("ready",)) - timeout_result = runtime.execute( - "import asyncio\n" - "await asyncio.gather(slow_tool(), slow_tool())\n", + result_queue.put(("ready",)) + result = runtime.execute( + "import asyncio\nawait asyncio.gather(slow_tool(), slow_tool())", timeout=0.1, ) - output = runtime.execute("print('still alive')") - queue.put(("ok", str(timeout_result), runtime.output(output))) + followup = runtime.execute("print('still alive')") + result_queue.put(("ok", result.timeout_seconds, runtime.output(followup))) except pytest.skip.Exception as exc: - queue.put(("skip", str(exc))) + result_queue.put(("skip", str(exc))) except BaseException as exc: - queue.put(("error", type(exc).__name__, str(exc))) + result_queue.put(("error", type(exc).__name__, str(exc))) finally: if runtime is not None: runtime.shutdown() -def _cleanup_real_sbx_runtime_sandbox(runtime: RuntimeHandle) -> None: - if runtime.spec.name != "sbx": - return - if os.environ.get("PREDICT_RLM_RUN_SBX_TESTS") != "1" or shutil.which("sbx") is None: - return - subprocess.run( - ["sbx", "rm", "-f", _REAL_SBX_RUNTIME_SANDBOX], - check=False, - capture_output=True, - text=True, - timeout=30, - ) - - -def _get_repro_message( - *, - process: multiprocessing.Process, - result_queue: multiprocessing.Queue, - timeout: float, - runtime: RuntimeHandle, - failure_message: str, -) -> tuple: +def _get_message(process, result_queue, timeout): deadline = time.monotonic() + timeout while time.monotonic() < deadline: - remaining = deadline - time.monotonic() try: - return result_queue.get(timeout=min(0.05, max(0.0, remaining))) + return result_queue.get(timeout=0.05) except queue_module.Empty: if not process.is_alive(): break - process.kill() - process.join(timeout=2) - _cleanup_real_sbx_runtime_sandbox(runtime) - pytest.fail(failure_message) - + pytest.fail("runtime stalled during concurrent host-tool timeout recovery") -def test_host_tool_round_trip(runtime: RuntimeHandle) -> None: - runtime.require("host_tools") - result = runtime.execute( - "result = await predict('question -> answer', question='2+2?')\n" - "print(result['answer'])" - ) - - assert runtime.output(result) == "4\n" - - -# Sized to greatly exceed the ~64KB OS pipe buffer while staying under the 1MB -# host-tool payload cap. Mirrors a real image `predict` tool-call request, which -# is hundreds of KB. -_LARGE_TOOL_REQUEST_BYTES = 950_000 -# The stall is intermittent per message, so repeat enough times that it is -# virtually certain to surface on the affected transport. -_LARGE_TOOL_REQUEST_ATTEMPTS = 12 - - -# TEMPORARY SKIP -- needs real investigation, do NOT just delete. -# This test was written for the OLD Docker `sbx exec` transport (a multiplexed pipe -# stream that wedged on >64KB inline messages); that transport is gone (real SBX is -# websocket now). On loaded CI runners the 950KB payload intermittently wedges the -# *pipe* backends and the request times out at 10s -- and not only the deprecated -# SbxBackend stdin/stdout seam, but also `python-runner/direct-process` -# (DirectPythonBackend, which uses the shared base pipe transport). It passes instantly -# everywhere locally (~0.3s), so it doesn't reproduce off-CI. -# OPEN QUESTION for the follow-up: is this a genuine large-message deadlock in the base -# pipe transport (reader-thread drain racing the inline read), or just CI slowness? If -# genuine, it's a real bug in DirectPythonBackend, not dead code -- which is why this is -# skipped-with-a-flag rather than removed. Re-enable once that's answered. -# local-only: CI-flaky large-message pipe wedge; needs investigation (see comment above) -- temporary -@pytest.mark.local -def test_large_host_tool_request_round_trips(runtime: RuntimeHandle) -> None: - """A large host-tool request must survive the host<->runner channel. - - The runner sends tool-call requests (e.g. an image `predict`) inline over the - host<->runner JSON-RPC channel. On the real Docker `sbx exec` backend that - channel is a multiplexed exec stream driven by blocking, unchunked pipe I/O, - and a single inline message far larger than the ~64KB pipe buffer - intermittently wedges it: the host never receives the request, the iteration - hits its watchdog, and the container-restart recovery fails ("container - started but not ready for exec"). The direct-pipe local and Deno backends - drain via separate pipes plus a reader thread, so they round-trip the same - payload fine. Repeated to make the intermittent stall reliably reproduce. - """ - runtime.require("host_tools") - - for attempt in range(_LARGE_TOOL_REQUEST_ATTEMPTS): - result = runtime.execute( - f"payload = 'x' * {_LARGE_TOOL_REQUEST_BYTES}\n" - "res = await predict('text: str -> answer: str', text=payload)\n" - "print(res['answer'])" - ) - assert runtime.output(result) == "4\n", ( - f"large host-tool request stalled on attempt {attempt}" - ) - - -def test_basic_host_tool_result_shapes(runtime: RuntimeHandle) -> None: - runtime.require("host_tools") - - result = runtime.execute( - "items = await shape_tool('list')\n" - "mapping = await shape_tool('dict')\n" - "none_value = await shape_tool('none')\n" - "text = await shape_tool('text')\n" - "print(items)\n" - "print(mapping['ok'])\n" - "print(none_value is None)\n" - "print(text)" - ) - - assert runtime.output(result) == "[1, 2]\nTrue\nTrue\nhello\n" - - -def test_recoverable_tool_exception_allows_later_tool_use(runtime: RuntimeHandle) -> None: - runtime.require("host_tools") - runtime.require("recoverable_errors") - - result = runtime.execute( - "try:\n" - " await failing_tool()\n" - "except Exception as exc:\n" - " print(type(exc).__name__)\n" - "print((await predict('question -> answer', question='2+2?'))['answer'])" - ) - - assert runtime.output(result).endswith("4\n") - - -def test_timeout_during_concurrent_host_tool_calls_is_recoverable( +def test_timeout_during_concurrent_host_tools_is_recoverable( runtime: RuntimeHandle, tmp_path: Path, ) -> None: - runtime.require("host_tools") - runtime.require("recoverable_iteration_timeout") - _cleanup_real_sbx_runtime_sandbox(runtime) - queue: multiprocessing.Queue = multiprocessing.Queue() + result_queue = multiprocessing.Queue() process = multiprocessing.Process( - target=_run_concurrent_host_tool_timeout_repro, - args=(runtime.spec.name, str(tmp_path / "staging"), queue), + target=_run_timeout_repro, + args=(runtime.spec.name, str(tmp_path / "staging"), result_queue), ) process.start() - status, *payload = _get_repro_message( - process=process, - result_queue=queue, - timeout=30, - runtime=runtime, - failure_message="SBX runtime did not finish startup before timeout recovery repro", - ) - if status == "skip": - pytest.skip(payload[0]) - assert status == "ready", payload - - status, *payload = _get_repro_message( - process=process, - result_queue=queue, - timeout=6, - runtime=runtime, - failure_message=( - "SBX supervisor hung after iteration timeout while awaiting " - "concurrent host tool calls" - ), - ) - - process.join(timeout=20) - if process.is_alive(): - process.kill() + try: + status, *payload = _get_message(process, result_queue, 30) + if status == "skip": + pytest.skip(payload[0]) + assert status == "ready", payload + status, *payload = _get_message(process, result_queue, 8) + assert status == "ok", payload + assert payload == [0.1, "still alive\n"] + process.join(timeout=20) + assert not process.is_alive(), "runtime cleanup stalled after recovery" + finally: + if process.is_alive(): + process.kill() process.join(timeout=2) - _cleanup_real_sbx_runtime_sandbox(runtime) - pytest.fail("SBX runtime cleanup hung after timeout recovery succeeded") - - assert status == "ok", payload - timeout_result, output = payload - assert "[Timeout] Iteration execution timed out after 0.1s" in timeout_result - assert output.strip() == "still alive" + result_queue.close() + result_queue.join_thread() + if ( + runtime.spec.name == "sbx" + and os.environ.get("PREDICT_RLM_RUN_SBX_TESTS") == "1" + and shutil.which("sbx") is not None + ): + subprocess.run( + ["sbx", "rm", "-f", "runtime-contract-sbx"], + check=False, + capture_output=True, + text=True, + timeout=30, + ) diff --git a/tests/test_adapter_contracts.py b/tests/test_adapter_contracts.py index 9dc8b350..d8b65a5b 100644 --- a/tests/test_adapter_contracts.py +++ b/tests/test_adapter_contracts.py @@ -2,19 +2,15 @@ from pathlib import Path from types import SimpleNamespace -from typing import Annotated import pytest -from pydantic import BaseModel from predict_rlm import ExecutionSpec, File, HostDirectoryMount from predict_rlm.compatibility import FileInputAdapter, FileOutputAdapter from predict_rlm.runtime import ( Artifact, - ArtifactBinding, FieldDescriptor, InputAdapter, - OutputAdapter, OutputReservation, PreparedInput, PreparedInputBinding, @@ -58,44 +54,6 @@ async def prepare(self, field, value, ctx): assert selected is concrete -@pytest.mark.parametrize( - ("annotation", "is_list", "allows_none", "item_allows_none"), - [ - (File, False, False, False), - (Annotated[File, "metadata"], False, False, False), - (File | None, False, True, False), - (Annotated[File | None, "metadata"], False, True, False), - (list[File], True, False, False), - (list[Annotated[File, "metadata"]], True, False, False), - (list[File] | None, True, True, False), - (list[File | None], True, False, True), - ], -) -def test_field_descriptor_normalizes_supported_annotation_shapes( - annotation, - is_list, - allows_none, - item_allows_none, -): - field = FieldDescriptor("source", annotation) - - assert field.name == "source" - assert field.matches(File) - assert field.is_list is is_list - assert field.allows_none is allows_none - assert field.item_allows_none is item_allows_none - - replacement = field.replace_type(str) - assert FieldDescriptor("replacement", replacement).is_list is is_list - assert FieldDescriptor("replacement", replacement).allows_none is allows_none - assert FieldDescriptor("replacement", replacement).item_allows_none is item_allows_none - - -def test_field_descriptor_does_not_flatten_arbitrary_unions_or_nested_lists(): - assert not FieldDescriptor("source", File | str).matches(File) - assert not FieldDescriptor("source", list[list[File]]).matches(File) - - def test_execution_spec_rejects_duplicate_host_mount_destinations(tmp_path): with pytest.raises(ValueError, match="Duplicate host-directory sandbox destination"): ExecutionSpec( @@ -116,28 +74,6 @@ def test_execution_spec_rejects_conflicting_host_mount_access(tmp_path): ) -def test_prepared_path_compiles_copy_without_adapter_plumbing(tmp_path): - source = tmp_path / "report.txt" - source.write_text("report", encoding="utf-8") - - prepared = compile_prepared_input( - FieldDescriptor("source", str), - PreparedInput.path(source), - ) - - assert prepared.model_value == "/sandbox/input/source/report.txt" - assert [dict(artifact.metadata) for artifact in prepared.artifacts] == [ - { - "source_path": str(source.resolve()), - "sandbox_path": "/sandbox/input/source/report.txt", - } - ] - assert prepared.requirements.extra_read_paths == (str(source.resolve()),) - assert [reservation.path for reservation in prepared.sandbox_roots] == [ - "/sandbox/input/source/report.txt" - ] - - def test_prepared_path_and_paths_honor_relative_destinations(tmp_path): first = tmp_path / "first.csv" second = tmp_path / "second.csv" @@ -177,9 +113,9 @@ def test_prepared_paths_reject_duplicate_destinations_within_one_field(tmp_path) ) with pytest.raises(ValueError, match="sandbox destinations overlap"): - validate_sandbox_root_reservations({ - field.name: PreparedInputBinding(field, FileInputAdapter(), prepared) - }) + validate_sandbox_root_reservations( + {field.name: PreparedInputBinding(field, FileInputAdapter(), prepared)} + ) def test_output_reservation_rejects_overlap_with_prepared_input(tmp_path): @@ -190,9 +126,7 @@ def test_output_reservation_rejects_overlap_with_prepared_input(tmp_path): field, PreparedInput.path(source, at="output/report"), ) - input_bindings = { - field.name: PreparedInputBinding(field, FileInputAdapter(), prepared) - } + input_bindings = {field.name: PreparedInputBinding(field, FileInputAdapter(), prepared)} output_field = FieldDescriptor("report", File) output = OutputReservation( field=output_field, @@ -208,72 +142,6 @@ def test_output_reservation_rejects_overlap_with_prepared_input(tmp_path): validate_output_sandbox_root_reservation(input_bindings, {}, output) -@pytest.mark.asyncio -async def test_pydantic_input_adapter_only_prepares_a_path(tmp_path): - source = tmp_path / "object.json" - source.write_text("{}", encoding="utf-8") - - class S3File(BaseModel): - uri: str - - class S3FileAdapter(InputAdapter[S3File]): - name = "s3-file" - value_type = S3File - - async def prepare(self, field, value, ctx): - return PreparedInput.path(source) - - class Session: - async def mount(self, artifact): - return ArtifactBinding(artifact.id, artifact.metadata["sandbox_path"]) - - field = FieldDescriptor("document", S3File) - adapter = S3FileAdapter() - prepared = compile_prepared_input( - field, - await adapter.prepare(field, S3File(uri="s3://bucket/object.json"), object()), - ) - bound = await adapter.bind(field, prepared, object(), Session()) - - assert bound.model_value == "/sandbox/input/document/object.json" - assert [binding.path for binding in bound.bindings] == [bound.model_value] - - -@pytest.mark.asyncio -async def test_prepared_directory_mount_uses_default_adapter_bind(tmp_path): - source = tmp_path / "dataset" - source.mkdir() - - class DatasetAdapter(InputAdapter[str]): - name = "dataset" - value_type = str - - async def prepare(self, field, value, ctx): - return PreparedInput.path(value, mode="mount", read_only=True) - - class Session: - async def mount_host_directory(self, mount): - assert mount.host_path == str(source.resolve()) - assert mount.sandbox_path == "/sandbox/input/dataset" - assert mount.read_only is True - return mount.sandbox_path - - field = FieldDescriptor("dataset", str) - adapter = DatasetAdapter() - prepared = compile_prepared_input( - field, - await adapter.prepare(field, str(source), object()), - ) - bound = await adapter.bind(field, prepared, object(), Session()) - bound_again = await adapter.bind(field, prepared, object(), Session()) - - assert bound.model_value == "/sandbox/input/dataset" - assert [binding.path for binding in bound.bindings] == [ - "/sandbox/input/dataset" - ] - assert bound.bindings[0].artifact_id != bound_again.bindings[0].artifact_id - - def test_prepared_glob_is_sorted_filtered_and_preserves_relative_paths(tmp_path): root = tmp_path / "dataset" (root / "nested").mkdir(parents=True) @@ -311,9 +179,9 @@ def test_prepared_glob_rejects_empty_matches_by_default(tmp_path): include="**/*.csv", allow_empty=True, ) - assert compile_prepared_input( - FieldDescriptor("files", list[str]), prepared - ).model_value == [] + assert ( + compile_prepared_input(FieldDescriptor("files", list[str]), prepared).model_value == [] + ) def test_prepared_glob_rejects_symlinks_outside_source_root(tmp_path): @@ -328,127 +196,7 @@ def test_prepared_glob_rejects_symlinks_outside_source_root(tmp_path): @pytest.mark.asyncio -async def test_typed_adapter_bases_own_matching_and_output_preparation(): - class StringInput(InputAdapter[str]): - name = "string" - value_type = str - - async def prepare(self, field, value, ctx): - return PreparedInput(model_value=f"{field.name}:{value}") - - class StringOutput(OutputAdapter[str]): - name = "string" - value_type = str - - async def reserve(self, field, value, ctx, session): - raise NotImplementedError - - async def materialize(self, reservation, submitted_value, ctx, session): - raise NotImplementedError - - field = FieldDescriptor("message", str) - input_adapter = StringInput() - prepared = await input_adapter.prepare(field, "hello", object()) - - assert input_adapter.supports(field, "hello") - assert await input_adapter.open(field, prepared, object(), object()) is None - bound = await input_adapter.bind(field, prepared, object(), object()) - assert bound.model_value == "message:hello" - assert await input_adapter.after_execution( - field, - prepared, - object(), - object(), - None, - RuntimeError("failed"), - ) is None - assert await input_adapter.finalize( - field, - prepared, - object(), - object(), - None, - ) is None - assert StringOutput().supports(field) - assert await StringOutput().prepare_session(field, None, object()) is None - - -@pytest.mark.asyncio -async def test_file_input_adapter_preserves_nullable_list_items(tmp_path): - source = tmp_path / "source.txt" - source.write_text("source", encoding="utf-8") - field = FieldDescriptor("sources", list[File | None]) - - prepared = await FileInputAdapter().prepare( - field, - [File(path=str(source)), None], - SimpleNamespace(state={}), - ) - - assert prepared.model_value == ["/sandbox/input/sources/source.txt", None] - assert len(prepared.artifacts) == 1 - - -@pytest.mark.asyncio -async def test_file_output_adapter_collects_every_generated_list_item(tmp_path): - destination = tmp_path / "results" - destination.mkdir() - (destination / "stale.txt").write_text("stale", encoding="utf-8") - - class Context: - state = { - "output_host_dirs": {"results": str(destination)}, - } - - def bind(self, binding): - return None - - class Session: - async def mount(self, artifact): - return ArtifactBinding(artifact.id, artifact.metadata["sandbox_path"]) - - async def collect(self, artifact): - target = Path(artifact.metadata["destination_path"]) - assert artifact.metadata["directory"] is True - (target / "first").mkdir(parents=True, exist_ok=True) - (target / "second").mkdir(parents=True, exist_ok=True) - (target / "first" / "result.txt").write_text("first", encoding="utf-8") - (target / "second" / "result.txt").write_text("second", encoding="utf-8") - return str(target) - - adapter = FileOutputAdapter() - field = FieldDescriptor("results", list[File]) - reservation = await adapter.reserve(field, None, Context(), Session()) - - result = await adapter.materialize( - reservation, - [ - File(path="/sandbox/output/results/first/result.txt"), - ], - Context(), - Session(), - ) - - assert [item.path for item in result] == [ - str(destination / "first" / "result.txt"), - str(destination / "second" / "result.txt"), - ] - - -@pytest.mark.asyncio -@pytest.mark.parametrize( - "submitted_value", - [ - None, - File(path=""), - File(path="relative.txt"), - File(path="/sandbox/output/result/missing.txt"), - ], -) -async def test_file_output_adapter_scalar_falls_back_to_reserved_directory( - tmp_path, - submitted_value, -): +async def test_file_output_adapter_scalar_falls_back_to_reserved_directory(tmp_path): destination = tmp_path / "result" class Session: @@ -474,7 +222,7 @@ async def collect(self, artifact): result = await FileOutputAdapter().materialize( reservation, - submitted_value, + File(path="/sandbox/output/result/missing.txt"), SimpleNamespace(), Session(), ) @@ -562,56 +310,7 @@ async def collect(self, artifact): @pytest.mark.asyncio -async def test_file_output_adapter_uses_remapped_reservation_root(tmp_path): - destination = tmp_path / "results" - - class Context: - state = {"output_host_dirs": {"results": str(destination)}} - - def bind(self, binding): - return None - - class Session: - async def mount(self, artifact): - return ArtifactBinding(artifact.id, "/workspace/results") - - async def collect(self, artifact): - target = Path(artifact.metadata["destination_path"]) - assert artifact.metadata["directory"] is True - (target / "nested").mkdir(parents=True, exist_ok=True) - (target / "nested" / "result.txt").write_text("result", encoding="utf-8") - return str(target) - - adapter = FileOutputAdapter() - reservation = await adapter.reserve( - FieldDescriptor("results", list[File]), - None, - Context(), - Session(), - ) - - result = await adapter.materialize( - reservation, - [File(path="/workspace/results/nested/result.txt")], - Context(), - Session(), - ) - - assert [item.path for item in result] == [str(destination / "nested" / "result.txt")] - - -@pytest.mark.asyncio -@pytest.mark.parametrize( - "submitted_path", - [ - "/sandbox/input/source.txt", - "/sandbox/output/results/../outside.txt", - ], -) -async def test_file_output_adapter_rejects_paths_outside_reservation( - tmp_path, - submitted_path, -): +async def test_file_output_adapter_rejects_paths_outside_reservation(tmp_path): artifact = SimpleNamespace( id="output", metadata={ @@ -627,7 +326,7 @@ async def test_file_output_adapter_rejects_paths_outside_reservation( with pytest.raises(ValueError, match="reserved output root"): await FileOutputAdapter().materialize( reservation, - [File(path=submitted_path)], + [File(path="/sandbox/output/results/../outside.txt")], SimpleNamespace(), SimpleNamespace(), ) diff --git a/tests/test_bootstrap_controller.py b/tests/test_bootstrap_controller.py index 48176d41..c35be9eb 100644 --- a/tests/test_bootstrap_controller.py +++ b/tests/test_bootstrap_controller.py @@ -8,40 +8,16 @@ import pytest REPO_ROOT = Path(__file__).resolve().parents[1] -BOOTSTRAP_SCRIPT = REPO_ROOT / "src/predict_rlm/remote/bootstrap_controller.sh" FIXTURE_ROOT = REPO_ROOT / "tests/fixtures/bootstrap_controller" DOCKER_SCENARIOS = ( "alpine", "busybox-unsupported-package-manager", - "debian13-no-python", - "python311-slim", "python313-slim-bookworm", - "tbench-python-313", "ubuntu24-no-python", "ubuntu24-nonroot-python-pip-no-venv", - "ubuntu24-python-pip-no-venv", ) -def test_bootstrap_controller_script_is_packaged_asset() -> None: - assert BOOTSTRAP_SCRIPT.is_file() - assert os.access(BOOTSTRAP_SCRIPT, os.X_OK) - assert BOOTSTRAP_SCRIPT.read_text(encoding="utf-8").startswith("#!/bin/sh\n") - - -def test_bootstrap_controller_script_is_valid_sh() -> None: - subprocess.run(["sh", "-n", str(BOOTSTRAP_SCRIPT)], check=True) - - -def test_bootstrap_controller_docker_matrix_covers_daytona_cases() -> None: - missing = [ - scenario - for scenario in DOCKER_SCENARIOS - if not (FIXTURE_ROOT / scenario / "Dockerfile").is_file() - ] - assert missing == [] - - @pytest.mark.parametrize("scenario", DOCKER_SCENARIOS) def test_bootstrap_controller_docker_scenario(scenario: str) -> None: if os.environ.get("PREDICT_RLM_RUN_BOOTSTRAP_DOCKER_TESTS") != "1": diff --git a/tests/test_callbacks.py b/tests/test_callbacks.py index 96e94f6f..3befdfe1 100644 --- a/tests/test_callbacks.py +++ b/tests/test_callbacks.py @@ -1,10 +1,4 @@ -"""Tests for PredictRLM lifecycle callbacks. - -Covers ``on_rlm_iteration_start`` and ``on_rlm_iteration_end`` handlers -dispatched from both the sync (``__call__``) and async (``acall``) -DSPy module paths. Uses a mocked interpreter and patched iteration helper -so no Deno sandbox is required — these tests run as pure unit tests. -""" +"""Callback failures, async handlers, and real iteration output.""" from __future__ import annotations @@ -17,9 +11,7 @@ from dspy.primitives.repl_types import REPLEntry, REPLHistory from dspy.utils.callback import BaseCallback -from predict_rlm import IterationStep, PredictRLM - -# --- Test signature & helpers --------------------------------------------- +from predict_rlm import PredictRLM class EchoSignature(dspy.Signature): @@ -52,7 +44,9 @@ def _final_prediction(answer: str = "ok") -> dspy.Prediction: pred = dspy.Prediction(answer=answer) # _execute_iteration's returned Prediction normally has a trajectory; the # loop reads .trajectory[-1] when it's a Prediction. Provide a minimal one. - pred.trajectory = [{"reasoning": "done", "code": "SUBMIT(answer='ok')", "output": "(no output)"}] + pred.trajectory = [ + {"reasoning": "done", "code": "SUBMIT(answer='ok')", "output": "(no output)"} + ] return pred @@ -72,24 +66,32 @@ def _build_rlm(**kwargs) -> PredictRLM: return rlm -def _drive_sync(rlm: PredictRLM, iteration_returns: list, fallback: dspy.Prediction | None = None): +def _drive_sync( + rlm: PredictRLM, iteration_returns: list, fallback: dspy.Prediction | None = None +): """Run rlm(...) with patched _execute_iteration returning the given sequence of values. ``fallback`` is used by _extract_fallback if we exhaust max_iterations without a final Prediction.""" fallback = fallback or _final_prediction(answer="fallback") - with patch.object(rlm, "_execute_iteration", side_effect=iteration_returns), \ - patch.object(rlm, "_extract_fallback", return_value=fallback), \ - dspy.context(lm=_make_lm()): + with ( + patch.object(rlm, "_execute_iteration", side_effect=iteration_returns), + patch.object(rlm, "_extract_fallback", return_value=fallback), + dspy.context(lm=_make_lm()), + ): return rlm(query="hi") -async def _drive_async(rlm: PredictRLM, iteration_returns: list, fallback: dspy.Prediction | None = None): +async def _drive_async( + rlm: PredictRLM, iteration_returns: list, fallback: dspy.Prediction | None = None +): fallback = fallback or _final_prediction(answer="fallback") aexec = AsyncMock(side_effect=iteration_returns) aextract = AsyncMock(return_value=fallback) - with patch.object(rlm, "_aexecute_iteration", aexec), \ - patch.object(rlm, "_aextract_fallback", aextract), \ - dspy.context(lm=_make_lm()): + with ( + patch.object(rlm, "_aexecute_iteration", aexec), + patch.object(rlm, "_aextract_fallback", aextract), + dspy.context(lm=_make_lm()), + ): return await rlm.acall(query="hi") @@ -103,22 +105,26 @@ def __init__(self): self.events: list[tuple[str, dict]] = [] def on_rlm_iteration_start(self, *, call_id, instance, iteration, max_iterations): - self.events.append(( - "start", - {"call_id": call_id, "iteration": iteration, "max_iterations": max_iterations}, - )) + self.events.append( + ( + "start", + {"call_id": call_id, "iteration": iteration, "max_iterations": max_iterations}, + ) + ) def on_rlm_iteration_end(self, *, call_id, instance, iteration, step, is_final, exception): - self.events.append(( - "end", - { - "call_id": call_id, - "iteration": iteration, - "step": step, - "is_final": is_final, - "exception": exception, - }, - )) + self.events.append( + ( + "end", + { + "call_id": call_id, + "iteration": iteration, + "step": step, + "is_final": is_final, + "exception": exception, + }, + ) + ) class AsyncRecordingCallback(BaseCallback): @@ -132,12 +138,16 @@ async def on_rlm_iteration_start(self, *, call_id, instance, iteration, max_iter await asyncio.sleep(0) self.events.append(("start", {"call_id": call_id, "iteration": iteration})) - async def on_rlm_iteration_end(self, *, call_id, instance, iteration, step, is_final, exception): + async def on_rlm_iteration_end( + self, *, call_id, instance, iteration, step, is_final, exception + ): await asyncio.sleep(0) - self.events.append(( - "end", - {"call_id": call_id, "iteration": iteration, "is_final": is_final}, - )) + self.events.append( + ( + "end", + {"call_id": call_id, "iteration": iteration, "is_final": is_final}, + ) + ) def _assert_shared_call_id(events: list[tuple[str, dict]]) -> None: @@ -146,69 +156,7 @@ def _assert_shared_call_id(events: list[tuple[str, dict]]) -> None: assert len(set(call_ids)) == 1 -# --- Sync path tests ------------------------------------------------------- - - class TestSyncCallbacks: - def test_events_fire_in_order_and_carry_step(self): - cb = RecordingCallback() - rlm = _build_rlm() - rlm.callbacks = [cb] - - h1 = _history_with(_entry(reasoning="r1", code="x=1", output="1")) - h2 = _history_with(_entry("r1", "x=1", "1"), _entry("r2", "x=2", "2")) - final = _final_prediction() - - _drive_sync(rlm, [h1, h2, final]) - - assert [name for name, _ in cb.events] == [ - "start", "end", "start", "end", "start", "end", - ] - _assert_shared_call_id(cb.events) - ends = [payload for name, payload in cb.events if name == "end"] - assert [e["iteration"] for e in ends] == [1, 2, 3] - assert [e["is_final"] for e in ends] == [False, False, True] - assert all(isinstance(e["step"], IterationStep) for e in ends) - assert ends[0]["step"].reasoning == "r1" - assert ends[1]["step"].code == "x=2" - # Final iteration's step is built from the Prediction's trajectory. - assert ends[2]["step"].reasoning == "done" - - def test_global_callback_via_dspy_settings(self): - cb = RecordingCallback() - rlm = _build_rlm() # no instance-level callback - with dspy.context(callbacks=[cb]): - _drive_sync(rlm, [_final_prediction()]) - assert [name for name, _ in cb.events] == ["start", "end"] - _assert_shared_call_id(cb.events) - - def test_instance_callback(self): - cb = RecordingCallback() - rlm = _build_rlm() - rlm.callbacks = [cb] - _drive_sync(rlm, [_final_prediction()]) - assert [name for name, _ in cb.events] == ["start", "end"] - _assert_shared_call_id(cb.events) - - def test_no_callbacks_runs_clean(self): - rlm = _build_rlm() - result = _drive_sync(rlm, [_final_prediction(answer="ok")]) - assert result.answer == "ok" - - def test_handler_exception_is_isolated(self, caplog): - class Boom(BaseCallback): - def on_rlm_iteration_start(self, **_): - raise RuntimeError("handler boom") - def on_rlm_iteration_end(self, **_): - raise RuntimeError("handler boom") - - rlm = _build_rlm() - rlm.callbacks = [Boom()] - with caplog.at_level(logging.WARNING, logger="predict_rlm.callbacks"): - result = _drive_sync(rlm, [_final_prediction(answer="ok")]) - assert result.answer == "ok" - assert any("handler boom" in rec.message for rec in caplog.records) - def test_iteration_end_fires_with_exception_when_iteration_raises(self): cb = RecordingCallback() rlm = _build_rlm() @@ -223,18 +171,6 @@ def test_iteration_end_fires_with_exception_when_iteration_raises(self): assert end_payload["is_final"] is False assert isinstance(end_payload["exception"], RuntimeError) - def test_max_iterations_status_emits_no_final_flag(self): - cb = RecordingCallback() - rlm = _build_rlm() - rlm.max_iterations = 2 - rlm.callbacks = [cb] - h = _history_with(_entry()) - _drive_sync(rlm, [h, _history_with(_entry(), _entry())]) - ends = [p for n, p in cb.events if n == "end"] - _assert_shared_call_id(cb.events) - assert all(e["is_final"] is False for e in ends) - assert len(ends) == 2 - def test_async_handler_in_sync_path_warns_and_skips(self, caplog): cb = AsyncRecordingCallback() rlm = _build_rlm() @@ -245,30 +181,8 @@ def test_async_handler_in_sync_path_warns_and_skips(self, caplog): assert cb.events == [] assert any("Async callback" in rec.message for rec in caplog.records) - def test_basecallback_subclass_without_rlm_methods_is_noop(self): - class OnlyLM(BaseCallback): - def on_lm_start(self, **_): pass - - rlm = _build_rlm() - rlm.callbacks = [OnlyLM()] - # Must not raise. - result = _drive_sync(rlm, [_final_prediction(answer="ok")]) - assert result.answer == "ok" - - -# --- Async path tests ------------------------------------------------------ - class TestAsyncCallbacks: - @pytest.mark.asyncio - async def test_sync_handler_in_async_path(self): - cb = RecordingCallback() - rlm = _build_rlm() - rlm.callbacks = [cb] - await _drive_async(rlm, [_final_prediction(answer="ok")]) - assert [name for name, _ in cb.events] == ["start", "end"] - _assert_shared_call_id(cb.events) - @pytest.mark.asyncio async def test_async_handler_is_awaited(self): cb = AsyncRecordingCallback() @@ -305,24 +219,7 @@ async def test_async_iteration_exception_still_emits_end(self): assert isinstance(cb.events[-1][1]["exception"], RuntimeError) -# --- Multi-callback ordering ---------------------------------------------- - - class TestMultipleCallbacks: - def test_global_and_instance_both_invoked(self): - global_cb = RecordingCallback() - instance_cb = RecordingCallback() - rlm = _build_rlm() - rlm.callbacks = [instance_cb] - with dspy.context(callbacks=[global_cb]): - _drive_sync(rlm, [_final_prediction()]) - # Both receive both events. - assert [n for n, _ in global_cb.events] == ["start", "end"] - assert [n for n, _ in instance_cb.events] == ["start", "end"] - _assert_shared_call_id(global_cb.events) - _assert_shared_call_id(instance_cb.events) - assert global_cb.events[0][1]["call_id"] == instance_cb.events[0][1]["call_id"] - def test_one_handler_failing_does_not_block_others(self): good = RecordingCallback() @@ -336,9 +233,6 @@ def on_rlm_iteration_start(self, **_): assert [n for n, _ in good.events] == ["start", "end"] -# --- Integration test (real Deno sandbox) --------------------------------- - - @pytest.mark.integration class TestCallbacksIntegration: """Verifies the callback contract end-to-end against a real interpreter. diff --git a/tests/test_codex_ws_lm.py b/tests/test_codex_ws_lm.py index 915a7479..50d4b985 100644 --- a/tests/test_codex_ws_lm.py +++ b/tests/test_codex_ws_lm.py @@ -23,117 +23,6 @@ def _events(): ] -@pytest.mark.asyncio -async def test_codex_wslm_default_transport_streams_responses_websocket(unused_tcp_port): - captured = {} - - async def ws_handler(request): - captured["headers"] = dict(request.headers) - ws = web.WebSocketResponse() - ws.headers["x-codex-turn-state"] = "sticky-1" - await ws.prepare(request) - message = await ws.receive_json() - captured["body"] = message - for event in _events(): - await ws.send_str(json.dumps(event)) - await ws.close() - return ws - - app = web.Application() - app.router.add_get("/backend-api/codex/responses", ws_handler) - runner = web.AppRunner(app) - await runner.setup() - port = unused_tcp_port - site = web.TCPSite(runner, "127.0.0.1", port) - await site.start() - - try: - lm = CodexWSLM( - model="gpt-5.3-codex", - access_token="fake-token", - account_id="fake-account", - ws_base="http://127.0.0.1:%d/backend-api/codex" % port, - cache=False, - ws_fallback=False, - ) - - result = await lm.aforward(prompt="hello", cache=False) - finally: - await runner.cleanup() - - assert result.output[0].content[0].text == "hello" - assert captured["body"]["type"] == "response.create" - assert captured["body"]["model"] == "gpt-5.3-codex" - assert captured["body"]["stream"] is True - assert captured["headers"]["OpenAI-Beta"] == "responses_websockets=2026-02-06" - assert captured["headers"]["Authorization"] == "Bearer fake-token" - assert captured["headers"]["ChatGPT-Account-Id"] == "fake-account" - - -@pytest.mark.asyncio -async def test_codex_wslm_moves_responses_lite_header_into_response_create_metadata( - unused_tcp_port, -): - captured = {"frames": []} - - async def ws_handler(request): - captured["headers"] = dict(request.headers) - ws = web.WebSocketResponse() - await ws.prepare(request) - captured["frames"].append(await ws.receive_json()) - await ws.send_json( - { - "type": "response.completed", - "response": {"id": "resp-prewarm"}, - } - ) - captured["frames"].append(await ws.receive_json()) - for event in _events(): - await ws.send_str(json.dumps(event)) - await ws.close() - return ws - - app = web.Application() - app.router.add_get("/backend-api/codex/responses", ws_handler) - runner = web.AppRunner(app) - await runner.setup() - port = unused_tcp_port - site = web.TCPSite(runner, "127.0.0.1", port) - await site.start() - - try: - lm = CodexWSLM( - model="gpt-5.6-sol", - access_token="fake-token", - account_id="fake-account", - ws_base="http://127.0.0.1:%d/backend-api/codex" % port, - cache=False, - ws_fallback=False, - ) - - await lm.aforward( - prompt="hello", - cache=False, - headers={"X-Test-Header": "keep"}, - client_metadata={"existing": "keep"}, - ) - finally: - await runner.cleanup() - - handshake_headers = {key.lower(): value for key, value in captured["headers"].items()} - assert "x-openai-internal-codex-responses-lite" not in handshake_headers - assert handshake_headers["x-test-header"] == "keep" - for frame in captured["frames"]: - assert frame["type"] == "response.create" - assert frame["client_metadata"]["existing"] == "keep" - assert ( - frame["client_metadata"][ - "ws_request_header_x_openai_internal_codex_responses_lite" - ] - == "true" - ) - - @pytest.mark.asyncio async def test_codex_wslm_response_lite_prewarms_before_generating(unused_tcp_port): frames = [] @@ -143,6 +32,11 @@ async def ws_handler(request): await ws.prepare(request) first = await ws.receive_json() + assert first["client_metadata"]["existing"] == "keep" + assert ( + first["client_metadata"]["ws_request_header_x_openai_internal_codex_responses_lite"] + == "true" + ) frames.append(first) if first.get("generate") is not False: for event in _events(): @@ -181,7 +75,9 @@ async def ws_handler(request): ws_fallback=False, ) - result = await lm.aforward(prompt="hello", cache=False) + result = await lm.aforward( + prompt="hello", cache=False, client_metadata={"existing": "keep"} + ) finally: await runner.cleanup() @@ -244,14 +140,14 @@ async def ws_handler(request): assert ( exc_info.value.failure_kind, exc_info.value.failure_code, - str(exc_info.value), exc_info.value.retry_after_seconds, ) == ( "error", "rate_limit_exceeded", - "Codex stream error (rate_limit_exceeded): prewarm capacity exhausted", 2.5, ) + assert "prewarm capacity exhausted" in str(exc_info.value) + @pytest.mark.asyncio async def test_codex_wslm_uses_fresh_turn_state_per_forward(unused_tcp_port): @@ -361,6 +257,7 @@ async def test_codex_wslm_websocket_401_is_codex_lm_auth_expired( unused_tcp_port, ): monkeypatch.setattr("dspy_codex_lm.lm.CODEX_STREAM_MAX_ATTEMPTS", 1) + async def auth_failed(_request): return web.Response(status=401, text="expired") diff --git a/tests/test_debug.py b/tests/test_debug.py deleted file mode 100644 index c16f48b8..00000000 --- a/tests/test_debug.py +++ /dev/null @@ -1,84 +0,0 @@ -import json - -import pytest - -from predict_rlm.debug import debug_event, reset_debug_logger_for_tests, sanitize_metadata - - -@pytest.fixture(autouse=True) -def reset_debug_logging(monkeypatch): - for name in ( - "PREDICT_RLM_DEBUG", - "RLM_DEBUG", - "PREDICT_RLM_DEBUG_LOG", - "PREDICT_RLM_DEBUG_JSON", - ): - monkeypatch.delenv(name, raising=False) - reset_debug_logger_for_tests() - yield - reset_debug_logger_for_tests() - - -def test_debug_logging_disabled_by_default(capsys): - debug_event("predict_rlm.test", count=1) - - captured = capsys.readouterr() - assert captured.err == "" - assert captured.out == "" - - -def test_predict_rlm_debug_enables_stderr_output(monkeypatch, capsys): - monkeypatch.setenv("PREDICT_RLM_DEBUG", "1") - - debug_event("predict_rlm.test", count=1) - - captured = capsys.readouterr() - assert "predict_rlm.test" in captured.err - assert "count=1" in captured.err - - -def test_rlm_debug_enables_stderr_output(monkeypatch, capsys): - monkeypatch.setenv("RLM_DEBUG", "1") - - debug_event("predict_rlm.shared", enabled=True) - - captured = capsys.readouterr() - assert "predict_rlm.shared" in captured.err - assert "enabled=True" in captured.err - - -def test_predict_rlm_debug_log_writes_file(monkeypatch, tmp_path, capsys): - log_path = tmp_path / "predict-rlm-debug.log" - monkeypatch.setenv("PREDICT_RLM_DEBUG", "1") - monkeypatch.setenv("PREDICT_RLM_DEBUG_LOG", str(log_path)) - - debug_event("predict_rlm.file", status="ok") - - captured = capsys.readouterr() - assert captured.err == "" - assert "predict_rlm.file" in log_path.read_text() - - -def test_json_debug_logging_redacts_obvious_secrets(monkeypatch, tmp_path): - log_path = tmp_path / "predict-rlm-debug.jsonl" - monkeypatch.setenv("PREDICT_RLM_DEBUG", "1") - monkeypatch.setenv("PREDICT_RLM_DEBUG_LOG", str(log_path)) - monkeypatch.setenv("PREDICT_RLM_DEBUG_JSON", "1") - - debug_event( - "predict_rlm.redact", - api_key="sk-testsecret123456", - nested={"authorization": "Bearer abcdefghijk"}, - harmless="visible", - ) - - record = json.loads(log_path.read_text()) - assert record["api_key"] == "[REDACTED]" - assert record["nested"]["authorization"] == "[REDACTED]" - assert record["harmless"] == "visible" - assert "sk-testsecret123456" not in log_path.read_text() - assert "Bearer abcdefghijk" not in log_path.read_text() - - -def test_sanitize_metadata_redacts_secret_values_in_nonsecret_keys(): - assert sanitize_metadata({"value": "Bearer abcdefghijk"}) == {"value": "[REDACTED]"} diff --git a/tests/test_direct_python_backend.py b/tests/test_direct_python_backend.py index 116da976..ebe5aae5 100644 --- a/tests/test_direct_python_backend.py +++ b/tests/test_direct_python_backend.py @@ -2,23 +2,9 @@ from pathlib import Path -from predict_rlm.backends.supervisor._payload import ( - _pickleable_globals_snapshot, - _SandboxPath, -) from predict_rlm.backends.supervisor.runner import DirectPythonBackend -def test_sandbox_path_snapshot_uses_virtual_string() -> None: - snapshot = _pickleable_globals_snapshot({ - "path": _SandboxPath("/sandbox/output/foo"), - }) - - assert snapshot["globals"] == {"path": "/sandbox/output/foo"} - assert snapshot["restored_globals"] == ["path"] - assert snapshot["lost_globals"] == [] - - def test_sandbox_path_global_does_not_poison_later_execute(tmp_path: Path) -> None: backend = DirectPythonBackend( runner_path=str(tmp_path / "predict_rlm_runner.py"), @@ -70,9 +56,7 @@ def test_regular_path_global_survives_timeout_recovery_as_path(tmp_path: Path) - ) try: first = backend.execute( - "from pathlib import Path\n" - "p = Path('/app/model.xml')\n" - "print(type(p).__name__)", + "from pathlib import Path\np = Path('/app/model.xml')\nprint(type(p).__name__)", timeout=10, ) backend.execute( @@ -83,9 +67,7 @@ def test_regular_path_global_survives_timeout_recovery_as_path(tmp_path: Path) - timeout=0.05, ) recovered = backend.execute( - "print(type(p).__name__)\n" - "print(hasattr(p, 'read_text'))\n" - "print(p / 'child.txt')", + "print(type(p).__name__)\nprint(hasattr(p, 'read_text'))\nprint(p / 'child.txt')", timeout=10, ) finally: diff --git a/tests/test_empty_code_retry.py b/tests/test_empty_code_retry.py index 996b8012..de15f9c6 100644 --- a/tests/test_empty_code_retry.py +++ b/tests/test_empty_code_retry.py @@ -1,232 +1,78 @@ -"""Regression coverage for invalid action ``code`` predictions. - -Background: - Some models return action-signature responses with empty or null ``code`` - fields. DSPy's adapter/Prediction boundary can otherwise let those values - materialize as successful ``Prediction`` objects even though the action - signature requires a non-empty string. - -Contract: - PredictRLM's validating adapter owns recovery: an empty ChatAdapter parse - should trigger JSON fallback, and JSON ``null`` should raise before a - ``Prediction`` is constructed. The RLM loop must not silently coerce invalid - predictions; if a malformed object reaches ``_aexecute_iteration`` directly - via a mock/custom adapter, it fails loudly. -""" - - -from __future__ import annotations +"""Invalid action outputs recover through JSON fallback or fail before execution.""" import asyncio -from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock import dspy import pytest - - -class _FakeRepl: - """Minimal repl with an async ``aexecute`` that just records the code.""" - - def __init__(self): - self.last_code = None - - async def aexecute(self, code, variables=None): - self.last_code = code - return "[Success] ok" - - -def _build_executor(): - """Make a bare PredictRLM instance we can call ``_aexecute_iteration`` on.""" - from predict_rlm.predict_rlm import PredictRLM - - executor = PredictRLM.__new__(PredictRLM) - executor.signature = dspy.Signature("question -> answer") - executor.max_iterations = 50 - executor.verbose = False - executor._user_tools = {} - # mock generate_action.acall so we control pred.code - executor.generate_action = MagicMock() - executor._partial_pending_entry = None - executor._partial_history = None - executor._process_execution_result = lambda pred, result, history, ofn: { - "result": result, - "pred_code": getattr(pred, "code", None), - } - return executor - - -def test_none_code_prediction_fails_loudly_in_rlm_loop(): - """A malformed direct Prediction means the validating adapter was bypassed.""" - executor = _build_executor() - - pred = SimpleNamespace(reasoning=None, code=None) - executor.generate_action.acall = AsyncMock(return_value=pred) - - repl = _FakeRepl() - - async def _run(): - return await executor._aexecute_iteration( - repl, - variables=[], - history=MagicMock(), - iteration=0, - input_args={}, - output_field_names=["answer"], +from dspy.primitives.repl_types import REPLHistory +from dspy.utils.exceptions import AdapterParseError + +from predict_rlm import PredictRLM +from predict_rlm.predict_rlm import _ValidatingChatAdapter + + +class ScriptedLM: + model = "openai/gpt-4o-mini" + supported_params = {"response_format", "temperature", "max_tokens"} + supports_response_schema = True + + def __init__(self, fallback): + self.fallback = fallback + self.responses = iter( + [ + "[[ ## reasoning ## ]]\ntry an action\n\n" + "[[ ## code ## ]]\n\n[[ ## completed ## ]]", + fallback, + ] ) - with pytest.raises(RuntimeError, match="invalid reasoning"): - asyncio.run(_run()) + def __call__(self, messages=None, **kwargs): + return [next(self.responses, self.fallback)] + async def acall(self, messages=None, **kwargs): + await asyncio.sleep(0) + return self(messages=messages, **kwargs) -def test_empty_string_code_prediction_fails_loudly_in_rlm_loop(): - """Empty code should be rejected by the adapter, not executed by the loop.""" - executor = _build_executor() - pred = SimpleNamespace(reasoning="…", code="") - executor.generate_action.acall = AsyncMock(return_value=pred) - - repl = _FakeRepl() - - async def _run(): - return await executor._aexecute_iteration( - repl, - variables=[], - history=MagicMock(), - iteration=0, - input_args={}, - output_field_names=["answer"], - ) - - with pytest.raises(RuntimeError, match="invalid code"): - asyncio.run(_run()) - - -def test_code_field_has_min_length_validator(): - """Source-anchor: the action-signature's ``code`` field must carry - a ``min_length=1`` constraint so PredictRLM's validating adapter rejects - empty code and triggers the ChatAdapter → JSONAdapter fallback retry. - If someone removes the constraint, we go back to accepting empty - responses silently. - """ - import dspy - - from predict_rlm._shared import build_rlm_signatures - - class _Base(dspy.Signature): - q: str = dspy.InputField() - answer: str = dspy.OutputField() - - action_sig, _ = build_rlm_signatures( - _Base, - instructions_template="", - user_tools={}, - format_tool_docs=lambda _: "", - ) - - code_field = action_sig.model_fields.get("code") - assert code_field is not None, "code field missing from action sig" - - # Pydantic min_length lives in the field's ``metadata`` list as a - # ``MinLen(min_length=1)`` constraint. Check it's there — more - # robust than instantiating the whole Signature (which requires - # populating several unrelated fields just to trigger validation). - has_min_length = any( - getattr(constraint, "min_length", None) == 1 - for constraint in code_field.metadata - ) - assert has_min_length, ( - f"code field metadata lacks min_length=1 constraint — empty/null " - f"code responses from the LM will parse silently. Field metadata: " - f"{code_field.metadata}" - ) - - -def test_validating_adapter_retries_empty_chat_code_via_json_fallback(): - """An empty parsed ChatAdapter code field must be treated as a parse - failure so DSPy's JSON fallback gets a chance to recover. - """ - import dspy - - from predict_rlm._shared import build_rlm_signatures - from predict_rlm.predict_rlm import _ValidatingChatAdapter - - class _Base(dspy.Signature): - q: str = dspy.InputField() - answer: str = dspy.OutputField() +def _adapter_inputs(): + return { + "lm_kwargs": {}, + "signature": PredictRLM("q -> answer").generate_action.signature, + "demos": [], + "inputs": {"variables_info": "", "repl_history": REPLHistory(), "iteration": "1/1"}, + } - action_sig, _ = build_rlm_signatures( - _Base, - instructions_template="", - user_tools={}, - format_tool_docs=lambda _: "", - ) - class _FakeLM: - model = "openai/gpt-4o-mini" - supported_params = {"response_format", "temperature", "max_tokens"} - supports_response_schema = True - - def __init__(self): - self.calls = [] - - def __call__(self, messages=None, **kwargs): - self.calls.append({"messages": messages, "kwargs": kwargs}) - if len(self.calls) == 1: - return [ - "[[ ## reasoning ## ]]\n" - "try an action\n\n" - "[[ ## code ## ]]\n\n" - "[[ ## completed ## ]]" - ] - return ['{"reasoning": "retry succeeded", "code": "print(1)"}'] - - lm = _FakeLM() +def test_empty_chat_action_recovers_through_json_fallback(): result = _ValidatingChatAdapter()( - lm, - lm_kwargs={}, - signature=action_sig, - demos=[], - inputs={ - "variables_info": "", - "repl_history": MagicMock(), - "iteration": "1/1", - }, + ScriptedLM('{"reasoning": "retry succeeded", "code": "print(1)"}'), + **_adapter_inputs(), ) + assert result[0]["code"] == "print(1)" + assert result[0]["reasoning"] == "retry succeeded" + + +@pytest.mark.asyncio +async def test_invalid_json_fallback_exhausts_recovery(): + with pytest.raises(AdapterParseError): + await asyncio.wait_for( + _ValidatingChatAdapter().acall( + ScriptedLM('{"reasoning": "still invalid", "code": null}'), + **_adapter_inputs(), + ), + timeout=3, + ) - assert result == [ - { - "reasoning": "retry succeeded", - "execution_timeout_seconds": None, - "code": "print(1)", - } - ] - assert len(lm.calls) == 2 - - -def test_validating_json_adapter_rejects_null_code(): - """JSON ``null`` is syntactically valid but invalid for required - non-optional signature fields, so the adapter must raise. - """ - import dspy - from dspy.utils.exceptions import AdapterParseError - - from predict_rlm._shared import build_rlm_signatures - from predict_rlm.predict_rlm import _ValidatingJSONAdapter - - class _Base(dspy.Signature): - q: str = dspy.InputField() - answer: str = dspy.OutputField() - - action_sig, _ = build_rlm_signatures( - _Base, - instructions_template="", - user_tools={}, - format_tool_docs=lambda _: "", - ) - with pytest.raises(AdapterParseError, match="cannot be null"): - _ValidatingJSONAdapter().parse( - action_sig, - '{"reasoning": "looks like json", "code": null}', - ) +@pytest.mark.asyncio +async def test_custom_predictor_empty_code_never_reaches_execution(): + rlm = PredictRLM("q -> answer") + rlm.generate_action = MagicMock() + rlm.generate_action.acall = AsyncMock( + return_value=dspy.Prediction(reasoning="attempt", code="") + ) + repl = MagicMock() + repl.aexecute = AsyncMock(side_effect=AssertionError("invalid code was executed")) + with pytest.raises(RuntimeError, match="invalid code"): + await rlm._aexecute_iteration(repl, [], REPLHistory(), 0, {}, ["answer"]) diff --git a/tests/test_exception_race.py b/tests/test_exception_race.py deleted file mode 100644 index bf162327..00000000 --- a/tests/test_exception_race.py +++ /dev/null @@ -1,518 +0,0 @@ -""" -Tests verifying tools survive long RLM traces (bug reproduction validated). - -BACKGROUND: -During long RLM traces with many iterations and parallel tool calls, tools -would sometimes become unavailable (NameError: 'predict' is not defined). - -BUG REPRODUCTION: -The bug was successfully reproduced by creating a version of runner.js WITHOUT -the following fixes: -1. No `_repl_tools` module persistence -2. No re-injection of tools before each execution -3. No `await responseReaderPromise` after code execution - -With the unfixed runner: -- Single iteration: PASSES -- Multiple iterations: HANGS on second iteration - -The hang occurs because `responseReader` is still blocked on `stdin.next()` -when the main loop tries to read the next code block. The `responseReader` -"steals" the code block (treating it as a tool response), causing deadlock. - -FIXES APPLIED (in runner.js): -1. Store tools in `_repl_tools` module in `sys.modules` (survives globals corruption) -2. Re-inject tools before each execution from `registeredTools` array -3. Always `await responseReaderPromise` after code execution to ensure clean handoff - -These tests verify the fixes work correctly by simulating the production pattern: -1. Many iterations (like RLM's 15-20 iterations) -2. Parallel tool calls via asyncio.gather() -3. Exceptions that are caught inside the sandbox code -4. Checks that tools remain available after errors -""" - -import pytest - -from predict_rlm.backends import JspiBackend - -pytestmark = pytest.mark.integration - - -class TestToolsExistInNamespace: - """Tests that explicitly check if tools exist in the sandbox namespace.""" - - def test_tools_in_globals_across_iterations(self): - """ - Explicitly check if tools exist in globals() across iterations. - This is the actual bug pattern - tools disappear from namespace. - """ - call_count = 0 - - async def predict(item: str) -> dict: - nonlocal call_count - call_count += 1 - return {"item": item, "count": call_count} - - interp = JspiBackend( - tools={"predict": predict}, - preinstall_packages=False, - ) - - try: - for i in range(20): - # Explicitly check if 'predict' exists in namespace BEFORE calling it - code = f""" -import sys - -# Check multiple places where tool might exist -in_globals = 'predict' in globals() -in_dir = 'predict' in dir() -in_repl_tools = '_repl_tools' in sys.modules and hasattr(sys.modules['_repl_tools'], 'predict') - -if not in_globals and not in_dir: - if in_repl_tools: - print(f"Iteration {i}: TOOLS_LOST_FROM_GLOBALS but recoverable from _repl_tools") - else: - print(f"Iteration {i}: TOOLS_COMPLETELY_LOST") - SUBMIT("TOOLS_LOST") -else: - # Tools exist, use them - result = await predict("item_{i}") - print(f"Iteration {i}: predict returned {{result}}") -""" - result = interp.execute(code) - output = str(result) - - if "TOOLS_COMPLETELY_LOST" in output: - pytest.fail(f"BUG REPRODUCED: Tools completely lost at iteration {i}!") - - if "TOOLS_LOST_FROM_GLOBALS" in output: - print(f"Note: Tools lost from globals but recovered at iteration {i}") - - print(f"All {call_count} iterations completed with tools intact") - - finally: - interp.shutdown() - - def test_tools_after_heavy_async_corruption(self): - """ - Try to corrupt globals with heavy async operations, then check tool existence. - """ - call_count = 0 - - async def predict(item: str) -> dict: - import asyncio - - nonlocal call_count - call_count += 1 - await asyncio.sleep(0.01) - return {"item": item} - - interp = JspiBackend( - tools={"predict": predict}, - preinstall_packages=False, - ) - - try: - for i in range(10): - # Heavy async work that might corrupt state - code = f""" -import asyncio -import sys - -# Do lots of async work -async def heavy_work(): - tasks = [predict(f"item_{{j}}") for j in range(5)] - results = await asyncio.gather(*tasks, return_exceptions=True) - - # Trigger some errors - try: - x = undefined_var_{i} - except NameError: - pass - - return results - -results = await heavy_work() -print(f"Iteration {i}: got {{len(results)}} results") - -# Now check if tools still exist -in_globals = 'predict' in globals() -in_repl_tools = '_repl_tools' in sys.modules and hasattr(sys.modules['_repl_tools'], 'predict') - -if not in_globals: - if in_repl_tools: - print(f"GLOBALS_CORRUPTED_BUT_RECOVERED") - else: - print(f"TOOLS_COMPLETELY_LOST") - SUBMIT("TOOLS_LOST") -""" - result = interp.execute(code) - output = str(result) - - if "TOOLS_COMPLETELY_LOST" in output: - pytest.fail( - f"BUG REPRODUCED: Tools lost after heavy async at iteration {i}!" - ) - - print(f"All iterations completed, {call_count} total predict calls") - - finally: - interp.shutdown() - - -class TestToolsSurviveExceptions: - """Tests that tools survive when exceptions are caught inside the sandbox.""" - - def test_tools_available_after_caught_exception(self): - """ - Reproduce: Tool call succeeds, then exception is caught, tools still work. - - This simulates the RLM pattern where exceptions are caught and handled. - """ - results = [] - - async def slow_tool(value: str) -> str: - results.append(f"slow_tool called with: {value}") - return f"processed: {value}" - - interp = JspiBackend( - tools={"slow_tool": slow_tool}, - preinstall_packages=False, - ) - - try: - # Iteration 1: Call tool, catch exception - code1 = """ -import asyncio - -# Call the tool successfully -result = await slow_tool("iteration1") -print(f"Got: {result}") - -# Raise and catch an exception -try: - raise ValueError("Model returned invalid response") -except ValueError as e: - print(f"Caught error: {e}") -""" - result1 = interp.execute(code1) - assert "Got: processed: iteration1" in str(result1) - assert "Caught error" in str(result1) - - # Verify tool was called - assert len(results) == 1 - - # Iteration 2: Tools should still work - code2 = """ -try: - result = await slow_tool("iteration2") - print(f"SUCCESS: {result}") -except NameError as e: - print(f"TOOLS_LOST: {e}") -""" - result2 = interp.execute(code2) - output = str(result2) - - if "TOOLS_LOST" in output: - pytest.fail( - f"BUG REPRODUCED: Tools were lost after exception! Output: {output}" - ) - - assert "SUCCESS" in output - assert len(results) == 2 - - finally: - interp.shutdown() - - def test_rapid_iterations_with_errors(self): - """ - Multiple rapid iterations with errors - stress test. - """ - call_count = 0 - - async def count_tool() -> int: - nonlocal call_count - call_count += 1 - return call_count - - interp = JspiBackend( - tools={"count_tool": count_tool}, - preinstall_packages=False, - ) - - try: - for i in range(10): - code = f""" -try: - result = await count_tool() - print(f"Iteration {i}: count = {{result}}") - # Raise and catch an error - raise RuntimeError("Iteration {i} error") -except RuntimeError as e: - print(f"Caught: {{e}}") -""" - result = interp.execute(code) - assert f"Iteration {i}" in str(result) - assert "Caught" in str(result) - - # Final check: tools should still work - final_code = """ -try: - result = await count_tool() - print(f"FINAL: count = {result}") -except NameError as e: - print(f"TOOLS_LOST: {e}") -""" - final_result = interp.execute(final_code) - output = str(final_result) - - if "TOOLS_LOST" in output: - pytest.fail(f"BUG REPRODUCED: Tools lost after {call_count} iterations!") - - assert "FINAL" in output - assert call_count == 11, f"Expected 11 calls, got {call_count}" - - finally: - interp.shutdown() - - def test_parallel_tool_calls_with_error_processing(self): - """ - Exception occurs while processing results from parallel tool calls. - This is the most likely scenario for the production bug. - """ - call_log = [] - - async def predict(item: str) -> dict: - import asyncio - - await asyncio.sleep(0.01) # Simulate async work - call_log.append(f"predict({item})") - return {"item": item, "value": f"extracted_{item}"} - - interp = JspiBackend( - tools={"predict": predict}, - preinstall_packages=False, - ) - - try: - # Pattern from RLM: parallel predict calls, then error processing - code1 = """ -import asyncio - -# Start multiple predict calls -tasks = [ - predict("item1"), - predict("item2"), - predict("item3"), -] - -# Gather them -results = await asyncio.gather(*tasks) -print(f"Got {len(results)} results") - -# Error processing the results (caught) -try: - x = results[0]["nonexistent_field"] # KeyError -except KeyError as e: - print(f"Caught KeyError: {e}") -""" - result1 = interp.execute(code1) - assert "Got 3 results" in str(result1) - assert "Caught KeyError" in str(result1) - - # Check that tools were called - assert len(call_log) == 3 - - # Now try next iteration - tools might be lost! - code2 = """ -try: - result = await predict("iteration2_item") - print(f"SUCCESS: {result}") -except NameError as e: - print(f"TOOLS_LOST: {e}") -""" - result2 = interp.execute(code2) - output = str(result2) - - if "TOOLS_LOST" in output: - pytest.fail("BUG REPRODUCED: Tools lost after parallel calls + exception!") - - assert "SUCCESS" in output - assert len(call_log) == 4 - - finally: - interp.shutdown() - - -class TestProductionScenarioReproduction: - """ - Reproduce the exact production scenario from dspy-playground.py - """ - - def test_rlm_extractor_pattern_20_iterations(self): - """ - Simulate the actual RLM extractor pattern: - 1. 20 iterations (like real RLM) - 2. Each iteration calls predict() and other tools - 3. Some iterations have processing errors (caught) - 4. Tools should survive all iterations - """ - iteration_count = 0 - - async def predict(extraction_type: str) -> dict: - nonlocal iteration_count - iteration_count += 1 - return { - "items": [{"name": f"item_{iteration_count}", "value": iteration_count * 10}] - } - - async def search(query: str) -> list: - return [{"page": 1, "text": f"Found: {query}"}] - - async def get_pages() -> list: - return [{"page": 1}, {"page": 2}] - - interp = JspiBackend( - tools={"predict": predict, "search": search, "get_pages": get_pages}, - preinstall_packages=False, - ) - - tools_lost_at = None - - try: - for i in range(20): - if i % 3 == 2: - # Error iteration - code = f""" -import asyncio - -# Call predict -result = await predict("iteration_{i}") -print(f"Iteration {i}: predict returned {{len(result.get('items', []))}} items") - -# Simulate processing error (caught) -try: - bad_value = result["items"][0]["nonexistent_field"] -except KeyError as e: - print(f"Caught KeyError: {{e}}") -""" - else: - # Normal iteration with multiple tool calls - code = f""" -import asyncio - -# Multiple tool calls -try: - results = await asyncio.gather( - predict("iteration_{i}"), - search("query_{i}"), - ) - print(f"Iteration {i}: predict={{results[0]}}, search={{results[1]}}") - - pages = await get_pages() - print(f"Pages: {{pages}}") -except NameError as e: - print(f"TOOLS_LOST at iteration {i}: {{e}}") -""" - - result = interp.execute(code) - output = str(result) - - # Check for tools lost - if "TOOLS_LOST" in output: - tools_lost_at = i - break - - if tools_lost_at is not None: - pytest.fail(f"BUG REPRODUCED: Tools lost at iteration {tools_lost_at}!") - - # Final verification - final_code = """ -try: - p = await predict("final") - s = await search("final") - g = await get_pages() - print(f"FINAL SUCCESS: predict={p}, search={s}, pages={g}") -except NameError as e: - print(f"TOOLS_LOST: {e}") -""" - final_result = interp.execute(final_code) - output = str(final_result) - - if "TOOLS_LOST" in output: - pytest.fail("BUG REPRODUCED: Tools lost after 20 iterations!") - - assert "FINAL SUCCESS" in output - print(f"All tools survived {iteration_count} predict calls across 20 iterations") - - finally: - interp.shutdown() - - def test_heavy_parallel_load(self): - """ - Heavy parallel load - many concurrent predict calls per iteration. - This stress tests the concurrent tool call handling. - """ - call_count = 0 - - async def predict(item: dict) -> dict: - import asyncio - - nonlocal call_count - call_count += 1 - await asyncio.sleep(0.01) # Simulate latency - return {"extracted": item.get("name", "unknown"), "count": call_count} - - interp = JspiBackend( - tools={"predict": predict}, - preinstall_packages=False, - ) - - try: - for iteration in range(5): - # Each iteration does 10 parallel predict calls - code = f""" -import asyncio - -items = [dict(name=f"item_{{i}}", iteration={iteration}) for i in range(10)] - -try: - results = await asyncio.gather(*[predict(item) for item in items]) - print(f"Iteration {iteration}: got {{len(results)}} results") - - # Processing might fail - if {iteration} == 2: - raise ValueError("Simulated processing error") -except ValueError as e: - print(f"Caught: {{e}}") -except NameError as e: - print(f"TOOLS_LOST at iteration {iteration}: {{e}}") -""" - result = interp.execute(code) - output = str(result) - - if "TOOLS_LOST" in output: - pytest.fail( - f"BUG REPRODUCED: Tools lost at iteration {iteration} with heavy parallel load!" - ) - - assert f"Iteration {iteration}" in output - - # Final check - final_code = """ -try: - result = await predict({"name": "final_test"}) - print(f"FINAL SUCCESS: {result}") -except NameError as e: - print(f"TOOLS_LOST: {e}") -""" - final_result = interp.execute(final_code) - assert "FINAL SUCCESS" in str(final_result) - - # 5 iterations * 10 calls + 1 final = 51 - assert call_count == 51, f"Expected 51 calls, got {call_count}" - - finally: - interp.shutdown() diff --git a/tests/test_external_input_adapter_contracts.py b/tests/test_external_input_adapter_contracts.py index 2f3149ce..d7544083 100644 --- a/tests/test_external_input_adapter_contracts.py +++ b/tests/test_external_input_adapter_contracts.py @@ -82,9 +82,7 @@ async def bind(self, field, prepared, ctx, session): path for path in prepared.state.staging_root.rglob("*") if path.is_file() ): relative = source.relative_to(prepared.state.staging_root).as_posix() - await session.transfer_file( - FileTransfer(str(source), f"/repository/{relative}") - ) + await session.transfer_file(FileTransfer(str(source), f"/repository/{relative}")) return BoundInput(model_value="/repository") async def after_execution(self, field, prepared, ctx, session, result, error): @@ -116,9 +114,7 @@ async def _flush(state: RepositoryLifecycle, session) -> None: if not isinstance(session, MutableDirectorySession): raise TypeError("mutable repositories require sync-back support") manifest = await session.inspect_directory("/repository") - current = { - relative for relative, info in manifest.items() if info.type == "file" - } + current = {relative for relative, info in manifest.items() if info.type == "file"} for relative in state.baseline - current: (state.staging_root / relative).unlink(missing_ok=True) for relative in current: @@ -211,7 +207,6 @@ async def test_one_stateless_adapter_handles_interleaved_real_jspi_runs(tmp_path "flushes": ["error", "success", "final"], "finalized": True, } - assert vars(adapter) == {} class ServiceAdapter(InputAdapter[str]): @@ -224,45 +219,10 @@ async def prepare(self, field, value, ctx): model_value=json.dumps( {"url": "https://snapshot.internal:8443", "snapshot": value} ), - requirements=SessionRequirements( - allowed_domains=("snapshot.internal:8443",) - ), + requirements=SessionRequirements(allowed_domains=("snapshot.internal:8443",)), ) -@pytest.mark.integration -@pytest.mark.skipif(shutil.which("deno") is None, reason="requires Deno") -def test_service_plain_value_and_requirement_reach_real_jspi(monkeypatch): - from predict_rlm.backends.jspi import execution as jspi_execution - - actual_backend = jspi_execution.JspiBackend - captured = {} - - def build_backend(**kwargs): - captured["allowed_domains"] = kwargs["allowed_domains"] - return actual_backend(**kwargs) - - monkeypatch.setattr(jspi_execution, "JspiBackend", build_backend) - rlm = PredictRLM( - "service: str -> answer: str", - lm=MagicMock(history=[]), - adapters=[ServiceAdapter()], - max_iterations=1, - verbose=False, - ) - rlm.generate_action.acall = AsyncMock( - return_value=dspy.Prediction( - reasoning="read the plain service descriptor", - code="import json\nSUBMIT(answer=json.loads(service)['snapshot'])", - ) - ) - - result = rlm(service="snapshot-42") - - assert result.answer == "snapshot-42" - assert captured["allowed_domains"] == ["snapshot.internal:8443"] - - @pytest.mark.sbx @pytest.mark.asyncio async def test_service_requirement_is_rejected_before_reused_or_pooled_sbx_acquisition(): @@ -339,9 +299,7 @@ async def bind(self, field, prepared, ctx, session): if not isinstance(session, HostDirectorySession): raise TypeError("read-only datasets require host-directory mounts") return BoundInput( - model_value=await session.mount_host_directory( - prepared.host_directory_mounts[0] - ) + model_value=await session.mount_host_directory(prepared.host_directory_mounts[0]) ) @@ -367,9 +325,7 @@ def test_owned_sbx_enforces_external_read_only_mount(tmp_path: Path): rlm = PredictRLM( DatasetSignature, lm=MagicMock(history=[]), - execution=SbxExecutionBackend( - config=SbxConfig(name=f"input-adapter-ro-{os.getpid()}") - ), + execution=SbxExecutionBackend(config=SbxConfig(name=f"input-adapter-ro-{os.getpid()}")), adapters=[ReadOnlyDatasetAdapter()], max_iterations=1, verbose=False, diff --git a/tests/test_file_sync.py b/tests/test_file_sync.py index 7d3d0abb..362b95b5 100644 --- a/tests/test_file_sync.py +++ b/tests/test_file_sync.py @@ -1,564 +1,134 @@ -"""Tests for SyncedFile type annotations and host-side file sync during tool calls. +"""SyncedFile transfer, writeback, ownership, and recovery in a real sandbox.""" -Unit tests verify type detection via get_synced_file_params. Integration tests -verify the full flow: sandbox code calls a SyncedFile-annotated tool -> framework -syncs the file from sandbox MEMFS to the host -> tool runs on host -> framework -mounts the modified file back into the sandbox. -""" - -import os +import shutil from pathlib import Path from typing import Annotated import pytest +from dspy.utils.dummies import DummyLM +from predict_rlm import PredictRLM from predict_rlm.backends import JspiBackend -from predict_rlm.files import SyncedFile, get_synced_file_params - -# ─── Unit tests ─────────────────────────────────────────────────────────────── - - -class TestGetSyncedFileParams: - """Unit tests for get_synced_file_params introspection.""" - - def test_no_annotations_returns_empty(self): - def my_tool(x: str) -> str: - return x - - assert get_synced_file_params(my_tool) == {} - - def test_single_synced_param(self): - def my_tool(file_path: Annotated[Path, SyncedFile()]) -> str: - return "ok" - - result = get_synced_file_params(my_tool) - assert "file_path" in result - assert result["file_path"] == SyncedFile() - - def test_multiple_synced_params(self): - def my_tool( - input_path: Annotated[Path, SyncedFile(writeback=False)], - output_path: Annotated[Path, SyncedFile()], - ) -> str: - return "ok" - - result = get_synced_file_params(my_tool) - assert len(result) == 2 - assert result["input_path"] == SyncedFile(writeback=False) - assert result["output_path"] == SyncedFile() - - def test_writeback_false_preserved(self): - def my_tool(ref: Annotated[Path, SyncedFile(writeback=False)]) -> str: - return "ok" - - assert get_synced_file_params(my_tool)["ref"].writeback is False - - def test_host_dir_preserved(self): - def my_tool( - f: Annotated[Path, SyncedFile(host_dir="/tmp/custom")], - ) -> str: - return "ok" - - assert get_synced_file_params(my_tool)["f"].host_dir == "/tmp/custom" - - def test_mixed_annotated_and_plain(self): - def my_tool( - synced: Annotated[Path, SyncedFile()], - plain: str, - count: int = 5, - ) -> str: - return "ok" - - result = get_synced_file_params(my_tool) - assert list(result.keys()) == ["synced"] - - def test_function_without_type_hints(self): - def my_tool(x, y): - return x - - assert get_synced_file_params(my_tool) == {} - - def test_frozen_dataclass(self): - sf = SyncedFile() - with pytest.raises(AttributeError): - sf.writeback = False # type: ignore[misc] - - -class TestToolDocFormatting: - """Test that SyncedFile-annotated tools render correctly in tool docs.""" - - def test_annotated_path_renders_as_str(self): - from predict_rlm._shared import format_tool_docs_full - - def my_tool( - workbook: Annotated[Path, SyncedFile()], - name: str, - ) -> str: - """Process a workbook.""" - return "ok" - - docs = format_tool_docs_full({"my_tool": my_tool}) - assert "workbook: str" in docs - assert "name: str" in docs - assert "SyncedFile" not in docs - assert "Annotated" not in docs - +from predict_rlm.files import SyncedFile -# ─── Integration tests ─────────────────────────────────────────────────────── +pytestmark = [ + pytest.mark.integration, + pytest.mark.skipif(shutil.which("deno") is None, reason="requires Deno"), +] -@pytest.mark.integration -class TestSyncedFileIntegration: - """Integration tests for SyncedFile-annotated tools running through the interpreter. +def test_binary_writeback_is_visible_to_later_calls_and_cleans_temporary_files(): + received = [] - These tests verify the full flow: sandbox writes a file, calls a - SyncedFile-annotated tool, the framework syncs the file to the host, the tool - modifies it, and the modified file is mounted back into the sandbox. - """ + def invert(path: Annotated[Path, SyncedFile()]) -> int: + path = Path(path) + received.append(path) + data = path.read_bytes() + path.write_bytes(bytes(byte ^ 0xFF for byte in data)) + return len(data) - def test_tool_receives_host_path_and_file_is_synced_back(self): - """A SyncedFile tool gets a real host path with the sandbox file's content, - and the modified file is mounted back into the sandbox.""" - received_paths = [] - - def modify_file( - file_path: Annotated[Path, SyncedFile()], - ) -> str: - received_paths.append(file_path) - with open(file_path, "r") as f: - content = f.read() - assert content == "hello from sandbox" - with open(file_path, "w") as f: - f.write("modified by host tool") - return "ok" - - interpreter = JspiBackend( - preinstall_packages=False, tools={"modify_file": modify_file} - ) - try: - output = interpreter.execute(""" -with open("/tmp/test_file.txt", "w") as f: - f.write("hello from sandbox") - -result = await modify_file(file_path="/tmp/test_file.txt") -print(f"tool returned: {result}") - -with open("/tmp/test_file.txt", "r") as f: - content = f.read() -print(f"content after tool: {content}") -""") - assert "tool returned: ok" in str(output) - assert "content after tool: modified by host tool" in str(output) - assert len(received_paths) == 1 - assert received_paths[0] != "/tmp/test_file.txt" - assert os.path.basename(received_paths[0]) == "test_file.txt" - finally: - interpreter.shutdown() - - def test_tool_with_positional_arg(self): - """SyncedFile works when the sandbox passes the path as a positional arg.""" - def read_file(file_path: Annotated[Path, SyncedFile(writeback=False)]) -> str: - with open(file_path, "r") as f: - return f.read().strip() - - interpreter = JspiBackend( - preinstall_packages=False, tools={"read_file": read_file} - ) - try: - output = interpreter.execute(""" -with open("/tmp/pos_test.txt", "w") as f: - f.write("positional arg test") - -result = await read_file("/tmp/pos_test.txt") -print(f"got: {result}") -""") - assert "got: positional arg test" in str(output) - finally: - interpreter.shutdown() - - def test_binary_file_roundtrip(self): - """SyncedFile handles binary files correctly.""" - def flip_bytes(file_path: Annotated[Path, SyncedFile()]) -> str: - with open(file_path, "rb") as f: - data = f.read() - with open(file_path, "wb") as f: - f.write(bytes(b ^ 0xFF for b in data)) - return "ok" - - interpreter = JspiBackend( - preinstall_packages=False, tools={"flip_bytes": flip_bytes} - ) - try: - output = interpreter.execute(""" -data = bytes([0, 1, 2, 3, 255]) -with open("/tmp/binary_test.bin", "wb") as f: - f.write(data) - -result = await flip_bytes(file_path="/tmp/binary_test.bin") - -with open("/tmp/binary_test.bin", "rb") as f: - result_data = f.read() - -expected = bytes([255, 254, 253, 252, 0]) -assert result_data == expected, f"Expected {list(expected)}, got {list(result_data)}" -print("binary roundtrip ok") -""") - assert "binary roundtrip ok" in str(output) - finally: - interpreter.shutdown() - - def test_tool_without_synced_file_unchanged(self): - """Tools without SyncedFile annotations are unaffected.""" - def synced_tool( - file_path: Annotated[Path, SyncedFile()], - ) -> str: - return "synced" - - def plain_tool(msg: str) -> str: - return f"echo: {msg}" - - interpreter = JspiBackend( - preinstall_packages=False, - tools={"synced_tool": synced_tool, "plain_tool": plain_tool}, - ) - try: - output = interpreter.execute(""" -result = await plain_tool(msg="hello") -print(result) -""") - assert "echo: hello" in str(output) - finally: - interpreter.shutdown() - - def test_synced_file_with_nonexistent_file(self): - """Tool call with a nonexistent sandbox file produces an error.""" - def read_it(file_path: Annotated[Path, SyncedFile(writeback=False)]) -> str: - with open(file_path) as f: - return f.read() - - interpreter = JspiBackend( - preinstall_packages=False, tools={"read_it": read_it} - ) - try: - output = interpreter.execute(""" + code = """ +from pathlib import Path +path = Path('/tmp/nested/data.bin') +path.parent.mkdir(parents=True) +data = bytes(range(256)) +path.write_bytes(data) +assert await invert(str(path)) == 256 +assert path.read_bytes() == bytes(byte ^ 0xFF for byte in data) +assert await invert(path=str(path)) == 256 +assert path.read_bytes() == data +SUBMIT(answer='roundtrip complete') +""" + rlm = PredictRLM( + "query -> answer", + tools={"invert": invert}, + lm=DummyLM([{"reasoning": "round trip", "code": code}]), + max_iterations=1, + ) + prediction = rlm(query="invert twice") + assert prediction.answer == "roundtrip complete" + assert len(received) == 2 + assert all(not path.parent.exists() for path in received) + + +def test_async_tool_failure_cleans_staging_and_allows_later_writeback(): + received = [] + + async def mutate(path: Annotated[Path, SyncedFile()], fail: bool) -> str: + path = Path(path) + received.append(path) + assert path.read_text(encoding="utf-8") == "original" + path.write_text("changed", encoding="utf-8") + if fail: + raise ValueError("intentional failure") + return "updated" + + interpreter = JspiBackend(preinstall_packages=False, tools={"mutate": mutate}) + try: + output = interpreter.execute(""" +from pathlib import Path +path = Path('/tmp/recovery.txt') +path.write_text('original') try: - result = await read_it(file_path="/tmp/nonexistent_file.txt") - print(f"unexpected: {result}") -except Exception as e: - print(f"error: {e}") -""") - assert "error:" in str(output) - finally: - interpreter.shutdown() - - def test_tool_can_grow_file(self): - """SyncedFile handles a tool that makes the file larger.""" - def append_data(file_path: Annotated[Path, SyncedFile()]) -> str: - with open(file_path, "a") as f: - f.write("\nextra line from host") - return "ok" - - interpreter = JspiBackend( - preinstall_packages=False, tools={"append_data": append_data} - ) - try: - output = interpreter.execute(""" -with open("/tmp/grow_test.txt", "w") as f: - f.write("original line") - -await append_data(file_path="/tmp/grow_test.txt") - -with open("/tmp/grow_test.txt", "r") as f: - lines = f.readlines() -print(f"line count: {len(lines)}") -print(f"last line: {lines[-1].strip()}") -""") - assert "line count: 2" in str(output) - assert "last line: extra line from host" in str(output) - finally: - interpreter.shutdown() - - def test_multiple_synced_file_calls(self): - """Multiple SyncedFile tool calls in the same execution work correctly.""" - call_count = [0] - - def increment_file(file_path: Annotated[Path, SyncedFile()]) -> str: - call_count[0] += 1 - with open(file_path, "r") as f: - val = int(f.read().strip()) - with open(file_path, "w") as f: - f.write(str(val + 1)) - return str(val + 1) - - interpreter = JspiBackend( - preinstall_packages=False, tools={"increment_file": increment_file} - ) - try: - output = interpreter.execute(""" -with open("/tmp/counter.txt", "w") as f: - f.write("0") - -r1 = await increment_file(file_path="/tmp/counter.txt") -r2 = await increment_file(file_path="/tmp/counter.txt") -r3 = await increment_file(file_path="/tmp/counter.txt") - -with open("/tmp/counter.txt", "r") as f: - final = f.read().strip() -print(f"results: {r1}, {r2}, {r3}") -print(f"final: {final}") + await mutate(str(path), fail=True) +except Exception as exc: + assert 'intentional failure' in str(exc) +else: + raise AssertionError('tool failure was swallowed') +assert path.read_text() == 'original' +assert await mutate(str(path), fail=False) == 'updated' +assert path.read_text() == 'changed' +print('recovered') +""") + assert output == "recovered\n" + assert len(received) == 2 + assert all(not path.parent.exists() for path in received) + finally: + interpreter.shutdown() + + +def test_read_only_sync_retains_custom_host_directory_without_changing_sandbox(tmp_path): + host_dir = tmp_path / "host" + + def mutate( + path: Annotated[Path, SyncedFile(writeback=False, host_dir=str(host_dir))], + ) -> str: + path = Path(path) + assert path.parent == host_dir + assert path.read_text(encoding="utf-8") == "original" + path.write_text("host-only", encoding="utf-8") + return "inspected" + + interpreter = JspiBackend(preinstall_packages=False, tools={"mutate": mutate}) + try: + output = interpreter.execute(""" +from pathlib import Path +path = Path('/tmp/readonly.txt') +path.write_text('original') +assert await mutate(str(path)) == 'inspected' +assert path.read_text() == 'original' +print('sandbox unchanged') """) - assert "results: 1, 2, 3" in str(output) - assert "final: 3" in str(output) - assert call_count[0] == 3 - finally: - interpreter.shutdown() + assert output == "sandbox unchanged\n" + assert (host_dir / "readonly.txt").read_text(encoding="utf-8") == "host-only" + finally: + interpreter.shutdown() - def test_synced_file_tool_error_still_works(self): - """If a SyncedFile tool raises an error, subsequent calls still work.""" - def maybe_fail( - file_path: Annotated[Path, SyncedFile()], - should_fail: bool = False, - ) -> str: - if should_fail: - raise ValueError("intentional failure") - with open(file_path, "r") as f: - return f.read().strip() - interpreter = JspiBackend( - preinstall_packages=False, tools={"maybe_fail": maybe_fail} - ) - try: - output = interpreter.execute(""" -with open("/tmp/err_test.txt", "w") as f: - f.write("test data") +def test_missing_synced_file_is_reported_as_an_error(): + def read_file(path: Annotated[Path, SyncedFile(writeback=False)]) -> str: + return Path(path).read_text(encoding="utf-8") + interpreter = JspiBackend(preinstall_packages=False, tools={"read_file": read_file}) + try: + output = interpreter.execute(""" try: - await maybe_fail(file_path="/tmp/err_test.txt", should_fail=True) -except Exception as e: - print(f"caught: {e}") - -result = await maybe_fail(file_path="/tmp/err_test.txt", should_fail=False) -print(f"after error: {result}") -""") - assert "caught:" in str(output) - assert "after error: test data" in str(output) - finally: - interpreter.shutdown() - - def test_mixed_sync_and_synced_file_tools(self): - """SyncedFile and plain tools work together in the same interpreter.""" - def transform_file(file_path: Annotated[Path, SyncedFile()]) -> str: - with open(file_path, "r") as f: - content = f.read() - with open(file_path, "w") as f: - f.write(content.upper()) - return "ok" - - def compute(x: int, y: int) -> int: - return x + y - - interpreter = JspiBackend( - preinstall_packages=False, - tools={"transform_file": transform_file, "compute": compute}, - ) - try: - output = interpreter.execute(""" -total = await compute(x=10, y=20) -print(f"total: {total}") - -with open("/tmp/mixed_test.txt", "w") as f: - f.write("hello world") - -await transform_file(file_path="/tmp/mixed_test.txt") - -with open("/tmp/mixed_test.txt", "r") as f: - result = f.read() -print(f"transformed: {result}") -""") - assert "total: 30" in str(output) - assert "transformed: HELLO WORLD" in str(output) - finally: - interpreter.shutdown() - - def test_synced_file_with_nested_directory(self): - """SyncedFile handles files in nested sandbox directories.""" - def stamp_file(file_path: Annotated[Path, SyncedFile()]) -> str: - with open(file_path, "a") as f: - f.write(" [stamped]") - return "ok" - - interpreter = JspiBackend( - preinstall_packages=False, tools={"stamp_file": stamp_file} - ) - try: - output = interpreter.execute(""" -import os -os.makedirs("/tmp/deep/nested/dir", exist_ok=True) - -with open("/tmp/deep/nested/dir/data.txt", "w") as f: - f.write("deep file") - -await stamp_file(file_path="/tmp/deep/nested/dir/data.txt") - -with open("/tmp/deep/nested/dir/data.txt", "r") as f: - result = f.read() -print(f"result: {result}") -""") - assert "result: deep file [stamped]" in str(output) - finally: - interpreter.shutdown() - - def test_async_synced_file_tool(self): - """SyncedFile works with async tool functions too.""" - import asyncio - - async def async_transform( - file_path: Annotated[Path, SyncedFile()], - ) -> str: - await asyncio.sleep(0.01) - with open(file_path, "r") as f: - content = f.read() - with open(file_path, "w") as f: - f.write(content[::-1]) - return "ok" - - interpreter = JspiBackend( - preinstall_packages=False, tools={"async_transform": async_transform} - ) - try: - output = interpreter.execute(""" -with open("/tmp/async_test.txt", "w") as f: - f.write("abcdef") - -await async_transform(file_path="/tmp/async_test.txt") - -with open("/tmp/async_test.txt", "r") as f: - result = f.read() -print(f"reversed: {result}") -""") - assert "reversed: fedcba" in str(output) - finally: - interpreter.shutdown() - - def test_temp_dir_cleanup(self): - """Temp directories created for file sync are cleaned up.""" - temp_dirs = [] - - def capture_dir(file_path: Annotated[Path, SyncedFile()]) -> str: - temp_dirs.append(os.path.dirname(file_path)) - return "ok" - - interpreter = JspiBackend( - preinstall_packages=False, tools={"capture_dir": capture_dir} - ) - try: - interpreter.execute(""" -with open("/tmp/cleanup_test.txt", "w") as f: - f.write("data") -await capture_dir(file_path="/tmp/cleanup_test.txt") -""") - assert len(temp_dirs) == 1 - assert not os.path.exists(temp_dirs[0]) - finally: - interpreter.shutdown() - - def test_synced_file_preserves_tool_return_value(self): - """The tool's return value is correctly passed back to the sandbox.""" - def analyze_file( - file_path: Annotated[Path, SyncedFile(writeback=False)], - ) -> dict: - with open(file_path, "r") as f: - content = f.read() - return {"length": len(content), "lines": content.count("\n") + 1} - - interpreter = JspiBackend( - preinstall_packages=False, tools={"analyze_file": analyze_file} - ) - try: - output = interpreter.execute(""" -with open("/tmp/analyze_test.txt", "w") as f: - f.write("line1\\nline2\\nline3") - -result = await analyze_file(file_path="/tmp/analyze_test.txt") -print(f"length: {result['length']}, lines: {result['lines']}") -""") - assert "length: 17, lines: 3" in str(output) - finally: - interpreter.shutdown() - - def test_writeback_false_skips_mount_after(self): - """With writeback=False, the tool can modify the host file but the sandbox - file remains unchanged.""" - def modify_but_readonly( - file_path: Annotated[Path, SyncedFile(writeback=False)], - ) -> str: - with open(file_path, "w") as f: - f.write("modified on host") - return "ok" - - interpreter = JspiBackend( - preinstall_packages=False, - tools={"modify_but_readonly": modify_but_readonly}, - ) - try: - output = interpreter.execute(""" -with open("/tmp/readonly_test.txt", "w") as f: - f.write("original content") - -await modify_but_readonly(file_path="/tmp/readonly_test.txt") - -with open("/tmp/readonly_test.txt", "r") as f: - content = f.read() -print(f"content: {content}") -""") - # Sandbox file should still have original content - assert "content: original content" in str(output) - finally: - interpreter.shutdown() - - def test_host_dir_uses_specified_directory(self, tmp_path): - """When host_dir is specified, the file is synced there instead of temp.""" - custom_dir = str(tmp_path / "custom_host") - received_paths = [] - - def check_dir( - file_path: Annotated[Path, SyncedFile(host_dir=custom_dir)], - ) -> str: - received_paths.append(str(file_path)) - return "ok" - - interpreter = JspiBackend( - preinstall_packages=False, tools={"check_dir": check_dir} - ) - try: - interpreter.execute(""" -with open("/tmp/hostdir_test.txt", "w") as f: - f.write("data") -await check_dir(file_path="/tmp/hostdir_test.txt") -""") - assert len(received_paths) == 1 - assert received_paths[0].startswith(custom_dir) - # Custom dir should NOT be cleaned up - assert os.path.exists(custom_dir) - finally: - interpreter.shutdown() - - def test_non_string_param_skipped(self): - """If a SyncedFile param receives a non-string value, it's skipped.""" - def flexible_tool( - file_path: Annotated[Path, SyncedFile()] = None, - data: str = "default", - ) -> str: - return f"path={file_path}, data={data}" - - interpreter = JspiBackend( - preinstall_packages=False, tools={"flexible_tool": flexible_tool} - ) - try: - output = interpreter.execute(""" -result = await flexible_tool(data="hello") -print(result) -""") - assert "path=None, data=hello" in str(output) - finally: - interpreter.shutdown() + await read_file('/tmp/nonexistent.txt') +except Exception: + print('missing file rejected') +else: + raise AssertionError('missing file was accepted') +""") + assert output == "missing file rejected\n" + finally: + interpreter.shutdown() diff --git a/tests/test_files.py b/tests/test_files.py index 8f4f43c6..9c9b512e 100644 --- a/tests/test_files.py +++ b/tests/test_files.py @@ -1,995 +1,98 @@ -"""Tests for declarative file I/O types.""" +"""Declarative file I/O through PredictRLM and a real sandbox.""" from __future__ import annotations -import os -import tempfile -from typing import Optional +import shutil +from pathlib import Path import dspy import pytest - -from predict_rlm.files import ( - File, - build_file_instructions, - build_file_plan, - is_file_type, - scan_file_fields, -) - -# -- Model creation tests -- - - -class TestFile: - def test_create_with_path(self): - f = File(path="/tmp/test.pdf") - assert f.path == "/tmp/test.pdf" - - def test_default_path_is_none(self): - f = File() - assert f.path is None - - def test_create_with_explicit_none(self): - f = File(path=None) - assert f.path is None - - -class TestFileFromDir: - def test_from_dir_walks_directory(self): - with tempfile.TemporaryDirectory() as tmpdir: - with open(os.path.join(tmpdir, "a.txt"), "w") as f: - f.write("a") - subdir = os.path.join(tmpdir, "sub") - os.makedirs(subdir) - with open(os.path.join(subdir, "b.txt"), "w") as f: - f.write("b") - - files = File.from_dir(tmpdir) - assert len(files) == 2 - paths = {f.path for f in files} - assert os.path.join(tmpdir, "a.txt") in paths - assert os.path.join(subdir, "b.txt") in paths - - def test_from_dir_empty_directory(self): - with tempfile.TemporaryDirectory() as tmpdir: - files = File.from_dir(tmpdir) - assert files == [] - - def test_from_dir_returns_file_instances(self): - with tempfile.TemporaryDirectory() as tmpdir: - with open(os.path.join(tmpdir, "test.txt"), "w") as f: - f.write("test") - files = File.from_dir(tmpdir) - assert all(isinstance(f, File) for f in files) - - -class TestDeprecatedAliases: - def test_local_file_is_file(self): - from predict_rlm.files import LocalFile - assert LocalFile is File - - def test_local_dir_is_file(self): - from predict_rlm.files import LocalDir - assert LocalDir is File - - def test_output_file_is_file(self): - from predict_rlm.files import OutputFile - assert OutputFile is File - - def test_output_dir_is_file(self): - from predict_rlm.files import OutputDir - assert OutputDir is File - - -# -- Type detection tests -- - - -class TestIsFileType: - def test_file(self): - assert is_file_type(File) is True - - def test_str_is_not_file(self): - assert is_file_type(str) is False - - def test_optional_file(self): - assert is_file_type(Optional[File]) is True - - def test_list_file(self): - assert is_file_type(list[File]) is True - - -# -- scan_file_fields tests -- - - -class TestScanFileFields: - def test_no_file_fields(self): - class Sig(dspy.Signature): - query: str = dspy.InputField() - answer: str = dspy.OutputField() - - inputs, outputs = scan_file_fields(Sig) - assert inputs == {} - assert outputs == {} - - def test_input_file_field(self): - class Sig(dspy.Signature): - source: File = dspy.InputField() - answer: str = dspy.OutputField() - - inputs, outputs = scan_file_fields(Sig) - assert inputs == {"source": "file"} - assert outputs == {} - - def test_output_file_field(self): - class Sig(dspy.Signature): - query: str = dspy.InputField() - result: File = dspy.OutputField() - - inputs, outputs = scan_file_fields(Sig) - assert inputs == {} - assert outputs == {"result": "file"} - - def test_list_file_input(self): - class Sig(dspy.Signature): - documents: list[File] = dspy.InputField() - answer: str = dspy.OutputField() - - inputs, outputs = scan_file_fields(Sig) - assert inputs == {"documents": "list_file"} - - def test_list_file_output(self): - class Sig(dspy.Signature): - query: str = dspy.InputField() - results: list[File] = dspy.OutputField() - - inputs, outputs = scan_file_fields(Sig) - assert outputs == {"results": "list_file"} - - def test_mixed_file_fields(self): - class Sig(dspy.Signature): - source: File = dspy.InputField() - docs: list[File] = dspy.InputField() - excel: File = dspy.OutputField() - pdfs: list[File] = dspy.OutputField() - - inputs, outputs = scan_file_fields(Sig) - assert inputs == {"source": "file", "docs": "list_file"} - assert outputs == {"excel": "file", "pdfs": "list_file"} - - -# -- build_file_instructions tests -- - - -class TestBuildFileInstructions: - def test_input_only(self): - result = build_file_instructions( - input_mounts={"source": "/sandbox/input/source/report.pdf"}, - output_dirs={}, - ) - assert "source" in result - assert "/sandbox/input/source/report.pdf" in result - assert "Output" not in result - - def test_output_only(self): - result = build_file_instructions( - input_mounts={}, - output_dirs={"result": "/sandbox/output/result/"}, - ) - assert "result" in result - assert "/sandbox/output/result/" in result - assert "Input" not in result - - def test_both(self): - result = build_file_instructions( - input_mounts={"source": "/sandbox/input/source/report.pdf"}, - output_dirs={"result": "/sandbox/output/result/"}, - ) - assert "source" in result - assert "result" in result - - def test_list_input(self): - result = build_file_instructions( - input_mounts={ - "docs": [ - "/sandbox/input/docs/file1.pdf", - "/sandbox/input/docs/file2.pdf", - ] - }, - output_dirs={}, - ) - assert "docs" in result - assert "file1.pdf" in result - - -# -- build_file_plan tests -- - - -class TestBuildFilePlan: - def test_returns_none_when_no_file_fields(self): - result = build_file_plan({}, {}, {}) - assert result is None - - def test_input_file_plan(self): - with tempfile.NamedTemporaryFile(suffix=".pdf", delete=False) as f: - f.write(b"fake pdf") - tmp_path = f.name - - try: - plan = build_file_plan( - input_args={"source": File(path=tmp_path)}, - input_file_fields={"source": "file"}, - output_file_fields={}, - ) - assert plan is not None - assert len(plan["mounts"]) == 1 - host_path, virtual_path = plan["mounts"][0] - assert host_path == tmp_path - assert virtual_path.startswith("/sandbox/input/source/") - assert tmp_path in plan["read_paths"] - finally: - os.unlink(tmp_path) - - def test_output_file_plan(self): - plan = build_file_plan( - input_args={}, - input_file_fields={}, - output_file_fields={"result": "file"}, - ) - assert plan is not None - assert "/sandbox/output/result" in plan["output_dirs"] - assert "result" in plan["output_field_map"] - assert plan["output_field_map"]["result"]["kind"] == "file" - - def test_output_with_custom_output_dir(self): - with tempfile.TemporaryDirectory() as tmpdir: - plan = build_file_plan( - input_args={}, - input_file_fields={}, - output_file_fields={"result": "file"}, - output_dir=tmpdir, - ) - assert plan["write_dir"] == tmpdir - assert plan["output_field_map"]["result"]["host_dir"] == os.path.join( - tmpdir, "result" - ) - - def test_instructions_generated(self): - with tempfile.NamedTemporaryFile(suffix=".pdf", delete=False) as f: - f.write(b"fake pdf") - tmp_path = f.name - - try: - plan = build_file_plan( - input_args={"source": File(path=tmp_path)}, - input_file_fields={"source": "file"}, - output_file_fields={"result": "file"}, - ) - assert "## Files" in plan["instructions"] - assert "source" in plan["instructions"] - assert "result" in plan["instructions"] - finally: - os.unlink(tmp_path) - - def test_list_file_input_plan(self): - with tempfile.NamedTemporaryFile(suffix=".pdf", delete=False) as f1: - f1.write(b"a") - p1 = f1.name - with tempfile.NamedTemporaryFile(suffix=".pdf", delete=False) as f2: - f2.write(b"b") - p2 = f2.name - - try: - plan = build_file_plan( - input_args={"docs": [File(path=p1), File(path=p2)]}, - input_file_fields={"docs": "list_file"}, - output_file_fields={}, - ) - assert plan is not None - assert len(plan["mounts"]) == 2 - virtual_paths = [vp for _, vp in plan["mounts"]] - assert all("/sandbox/input/docs/" in vp for vp in virtual_paths) - finally: - os.unlink(p1) - os.unlink(p2) - - -# -- PredictRLM-level unit tests -- - - -class TestPrepareFileIO: - """Tests for PredictRLM._prepare_file_io.""" - - def _make_rlm(self, sig, **kwargs): - from unittest.mock import MagicMock - - from predict_rlm import PredictRLM - - return PredictRLM(sig, sub_lm=MagicMock(), max_iterations=1, **kwargs) - - def test_no_file_fields_returns_none(self): - class Sig(dspy.Signature): - query: str = dspy.InputField() - answer: str = dspy.OutputField() - - rlm = self._make_rlm(Sig) - plan, args = rlm._prepare_file_io({"query": "hello"}) - assert plan is None - assert args == {"query": "hello"} - - def test_input_file_transformed_to_path_string(self): - class Sig(dspy.Signature): - source: File = dspy.InputField() - answer: str = dspy.OutputField() - - with tempfile.NamedTemporaryFile(suffix=".pdf", delete=False) as f: - f.write(b"data") - tmp = f.name - - try: - rlm = self._make_rlm(Sig) - plan, args = rlm._prepare_file_io( - {"source": File(path=tmp)} - ) - assert plan is not None - basename = os.path.basename(tmp) - assert args["source"] == f"/sandbox/input/source/{basename}" - finally: - os.unlink(tmp) - - def test_output_fields_removed_from_args(self): - class Sig(dspy.Signature): - query: str = dspy.InputField() - result: File = dspy.OutputField() - - rlm = self._make_rlm(Sig) - plan, args = rlm._prepare_file_io( - {"query": "hello", "result": File()} - ) - assert "result" not in args - assert args == {"query": "hello"} - - def test_non_file_fields_preserved(self): - class Sig(dspy.Signature): - source: File = dspy.InputField() - query: str = dspy.InputField() - result: File = dspy.OutputField() - summary: str = dspy.OutputField() - - with tempfile.NamedTemporaryFile(suffix=".pdf", delete=False) as f: - f.write(b"data") - tmp = f.name - - try: - rlm = self._make_rlm(Sig) - plan, args = rlm._prepare_file_io({ - "source": File(path=tmp), - "query": "summarize", - "result": File(), - }) - assert "query" in args - assert args["query"] == "summarize" - assert "result" not in args - assert args["source"].startswith("/sandbox/input/") - finally: - os.unlink(tmp) - - def test_list_file_transformed_to_list_of_paths(self): - class Sig(dspy.Signature): - documents: list[File] = dspy.InputField() - answer: str = dspy.OutputField() - - with tempfile.NamedTemporaryFile(suffix=".pdf", delete=False) as f1: - f1.write(b"a") - p1 = f1.name - with tempfile.NamedTemporaryFile(suffix=".pdf", delete=False) as f2: - f2.write(b"b") - p2 = f2.name - - try: - rlm = self._make_rlm(Sig) - plan, args = rlm._prepare_file_io({ - "documents": [File(path=p1), File(path=p2)] - }) - assert plan is not None - assert isinstance(args["documents"], list) - assert len(args["documents"]) == 2 - assert all(p.startswith("/sandbox/input/documents/") for p in args["documents"]) - assert len(plan["mounts"]) == 2 - finally: - os.unlink(p1) - os.unlink(p2) - - -class TestBuildSignaturesWithFiles: - """Tests for PredictRLM._build_signatures_with_files.""" - - def _make_rlm(self, sig): - from unittest.mock import MagicMock - - from predict_rlm import PredictRLM - - return PredictRLM(sig, sub_lm=MagicMock(), max_iterations=1) - - def test_output_file_type_replaced_with_str(self): - class Sig(dspy.Signature): - query: str = dspy.InputField() - result: File = dspy.OutputField(desc="Generated file") - - rlm = self._make_rlm(Sig) - action, extract = rlm._build_signatures_with_files("## Files\ntest") - - assert "result" in action.signature.instructions - assert "SUBMIT(result)" in action.signature.instructions - - def test_file_instructions_in_action_signature(self): - class Sig(dspy.Signature): - query: str = dspy.InputField() - answer: str = dspy.OutputField() - - rlm = self._make_rlm(Sig) - file_instr = "## Files\n\n- source: /sandbox/input/source/test.pdf" - action, _ = rlm._build_signatures_with_files(file_instr) - assert "/sandbox/input/source/test.pdf" in action.signature.instructions - - def test_input_file_type_replaced_with_str(self): - class Sig(dspy.Signature): - source: File = dspy.InputField(desc="Input PDF") - answer: str = dspy.OutputField() - - rlm = self._make_rlm(Sig) - action, _ = rlm._build_signatures_with_files("## Files\ntest") - assert "`source`" in action.signature.instructions - - -class TestSyncOutputFiles: - """Tests for PredictRLM._sync_output_files.""" - - def _make_rlm(self, sig): - from unittest.mock import MagicMock - - from predict_rlm import PredictRLM - - return PredictRLM(sig, sub_lm=MagicMock(), max_iterations=1) - - def test_sync_with_valid_sandbox_path(self): - from unittest.mock import MagicMock - - class Sig(dspy.Signature): - query: str = dspy.InputField() - result: File = dspy.OutputField() - - rlm = self._make_rlm(Sig) - - mock_repl = MagicMock() - prediction = dspy.Prediction( - result="/sandbox/output/result/out.xlsx", query="test" - ) - - with tempfile.TemporaryDirectory() as tmpdir: - file_plan = { - "output_field_map": { - "result": { - "virtual_dir": "/sandbox/output/result", - "host_dir": os.path.join(tmpdir, "result"), - "kind": "file", - } - } - } - - rlm._sync_output_files( - mock_repl, prediction, {"result": "file"}, file_plan - ) - - mock_repl.sync_file_to.assert_called_once_with( - "/sandbox/output/result/out.xlsx", - os.path.join(tmpdir, "result", "out.xlsx"), - ) - assert isinstance(prediction.result, File) - assert prediction.result.path == os.path.join( - tmpdir, "result", "out.xlsx" - ) - - def test_sync_fallback_when_no_valid_path(self): - from unittest.mock import MagicMock - - class Sig(dspy.Signature): - query: str = dspy.InputField() - result: File = dspy.OutputField() - - rlm = self._make_rlm(Sig) - - mock_repl = MagicMock() - mock_repl.list_dir.return_value = ["/sandbox/output/result/file.csv"] - - prediction = dspy.Prediction( - result="some random text", query="test" - ) - - with tempfile.TemporaryDirectory() as tmpdir: - file_plan = { - "output_field_map": { - "result": { - "virtual_dir": "/sandbox/output/result", - "host_dir": os.path.join(tmpdir, "result"), - "kind": "file", - } +from dspy.utils.dummies import DummyLM + +from predict_rlm import File, PredictRLM + +pytestmark = [ + pytest.mark.integration, + pytest.mark.skipif(shutil.which("deno") is None, reason="requires Deno"), +] + + +def test_file_roundtrip_preserves_binary_data_and_hides_host_destination(tmp_path): + class Signature(dspy.Signature): + source: File = dspy.InputField() + result: File = dspy.OutputField() + + source = tmp_path / "source.bin" + contents = bytes(range(256)) + source.write_bytes(contents) + destination = tmp_path / "destination" + rlm = PredictRLM( + Signature, + lm=DummyLM( + [ + { + "reasoning": "transform the mounted file without seeing a host output path", + "code": ( + "from pathlib import Path\n" + "assert 'result' not in globals()\n" + "data = Path(source).read_bytes()\n" + "output = Path('/sandbox/output/result/transformed.bin')\n" + "output.write_bytes(data[::-1])\n" + "SUBMIT(result=str(output))" + ), } - } - - rlm._sync_output_files( - mock_repl, prediction, {"result": "file"}, file_plan - ) - - mock_repl.list_dir.assert_called_once_with("/sandbox/output/result") - mock_repl.sync_file_to.assert_called_once() - assert isinstance(prediction.result, File) - - def test_sync_list_file_output(self): - from unittest.mock import MagicMock - - class Sig(dspy.Signature): - query: str = dspy.InputField() - outfiles: list[File] = dspy.OutputField() - - rlm = self._make_rlm(Sig) - - mock_repl = MagicMock() - mock_repl.list_dir.return_value = [ - "/sandbox/output/outfiles/a.txt", - "/sandbox/output/outfiles/sub/b.txt", - ] - - prediction = dspy.Prediction( - outfiles="/sandbox/output/outfiles", query="test" - ) - - with tempfile.TemporaryDirectory() as tmpdir: - file_plan = { - "output_field_map": { - "outfiles": { - "virtual_dir": "/sandbox/output/outfiles", - "host_dir": os.path.join(tmpdir, "outfiles"), - "kind": "list_file", - } + ] + ), + max_iterations=1, + ) + + prediction = rlm(source=File(path=str(source)), result=File(path=str(destination))) + + assert Path(prediction.result.path) == destination / "transformed.bin" + assert Path(prediction.result.path).read_bytes() == contents[::-1] + assert source.read_bytes() == contents + + +def test_file_list_collects_generated_outputs_without_stale_host_files(tmp_path): + class Signature(dspy.Signature): + source: list[File] = dspy.InputField() + results: list[File] = dspy.OutputField() + + source = tmp_path / "source" + (source / "nested").mkdir(parents=True) + (source / "first.txt").write_text("first", encoding="utf-8") + (source / "nested" / "second.txt").write_text("second", encoding="utf-8") + output_dir = tmp_path / "outputs" + (output_dir / "results").mkdir(parents=True) + (output_dir / "results" / "stale.txt").write_text("stale", encoding="utf-8") + rlm = PredictRLM( + Signature, + lm=DummyLM( + [ + { + "reasoning": "retain nested paths and discover an unsubmitted output", + "code": ( + "from pathlib import Path\n" + "inputs = {Path(path).name: Path(path) for path in source}\n" + "root = Path('/sandbox/output/results')\n" + "for name, input_name in [('first', 'first.txt'), ('second', 'second.txt')]:\n" + " output = root / name / 'result.txt'\n" + " output.parent.mkdir(parents=True, exist_ok=True)\n" + " output.write_text(inputs[input_name].read_text().upper())\n" + "SUBMIT(results=[str(root / 'first/result.txt')])" + ), } - } - - rlm._sync_output_files( - mock_repl, prediction, {"outfiles": "list_file"}, file_plan - ) - - assert mock_repl.sync_file_to.call_count == 2 - assert isinstance(prediction.outfiles, list) - assert len(prediction.outfiles) == 2 - assert all(isinstance(f, File) for f in prediction.outfiles) - - def test_sync_no_files_written(self): - from unittest.mock import MagicMock - - class Sig(dspy.Signature): - query: str = dspy.InputField() - result: File = dspy.OutputField() - - rlm = self._make_rlm(Sig) - - mock_repl = MagicMock() - mock_repl.list_dir.return_value = [] - - prediction = dspy.Prediction(result="", query="test") - - file_plan = { - "output_field_map": { - "result": { - "virtual_dir": "/sandbox/output/result", - "host_dir": "/tmp/test-result", - "kind": "file", - } - } - } - - rlm._sync_output_files( - mock_repl, prediction, {"result": "file"}, file_plan - ) - - mock_repl.sync_file_to.assert_not_called() - - -class TestOutputFieldsInfo: - """Tests for _get_output_fields_info with File-typed fields.""" - - def test_file_output_field_gets_str_type(self): - """File output fields should appear as 'str' in SUBMIT signature.""" - from unittest.mock import MagicMock - - from predict_rlm import PredictRLM - - class Sig(dspy.Signature): - query: str = dspy.InputField() - workbook: File = dspy.OutputField(desc="output file") - result: str = dspy.OutputField(desc="summary") - - rlm = PredictRLM(Sig, sub_lm=MagicMock(), max_iterations=1) - info = rlm._get_output_fields_info() - - assert info == [ - {"name": "workbook", "type": "str"}, - {"name": "result", "type": "str"}, - ] - - def test_list_file_output_field_gets_list_type(self): - """list[File] output fields should appear as 'list' in SUBMIT signature.""" - from unittest.mock import MagicMock - - from predict_rlm import PredictRLM - - class Sig(dspy.Signature): - docs: list[File] = dspy.InputField() - redacted: list[File] = dspy.OutputField(desc="redacted files") - result: dict = dspy.OutputField(desc="summary") - - rlm = PredictRLM(Sig, sub_lm=MagicMock(), max_iterations=1) - info = rlm._get_output_fields_info() - - assert info == [ - {"name": "redacted", "type": "list"}, - {"name": "result", "type": "dict"}, - ] - - -class TestProcessFinalOutput: - """Tests for _process_final_output coercing strings to File for file fields.""" - - def test_string_path_coerced_to_file(self): - """A plain string for a File output field should be wrapped as {"path": str}.""" - from unittest.mock import MagicMock - - from dspy.primitives.code_interpreter import FinalOutput - - from predict_rlm import PredictRLM - - class Sig(dspy.Signature): - query: str = dspy.InputField() - workbook: File = dspy.OutputField(desc="output file") - result: str = dspy.OutputField(desc="summary") - - rlm = PredictRLM(Sig, sub_lm=MagicMock(), max_iterations=1) - final = FinalOutput({ - "workbook": "/sandbox/output/workbook/result.xlsx", - "result": "done", - }) - parsed, error = rlm._process_final_output(final, ["workbook", "result"]) - assert error is None, f"Unexpected error: {error}" - assert isinstance(parsed["workbook"], File) - assert parsed["workbook"].path == "/sandbox/output/workbook/result.xlsx" - assert parsed["result"] == "done" - - def test_list_string_paths_coerced_to_files(self): - """A list of strings for a list[File] field should each be wrapped.""" - from unittest.mock import MagicMock - - from dspy.primitives.code_interpreter import FinalOutput - - from predict_rlm import PredictRLM - - class Sig(dspy.Signature): - docs: list[File] = dspy.InputField() - redacted: list[File] = dspy.OutputField(desc="redacted files") - - rlm = PredictRLM(Sig, sub_lm=MagicMock(), max_iterations=1) - final = FinalOutput({ - "redacted": ["/sandbox/output/redacted/a.pdf", "/sandbox/output/redacted/b.pdf"], - }) - parsed, error = rlm._process_final_output(final, ["redacted"]) - assert error is None, f"Unexpected error: {error}" - assert len(parsed["redacted"]) == 2 - assert all(isinstance(f, File) for f in parsed["redacted"]) - assert parsed["redacted"][0].path == "/sandbox/output/redacted/a.pdf" - - def test_dict_path_still_works(self): - """A dict {"path": ...} for a File field should still work (no double-wrapping).""" - from unittest.mock import MagicMock - - from dspy.primitives.code_interpreter import FinalOutput - - from predict_rlm import PredictRLM - - class Sig(dspy.Signature): - query: str = dspy.InputField() - workbook: File = dspy.OutputField(desc="output file") - - rlm = PredictRLM(Sig, sub_lm=MagicMock(), max_iterations=1) - final = FinalOutput({ - "workbook": {"path": "/sandbox/output/workbook/result.xlsx"}, - }) - parsed, error = rlm._process_final_output(final, ["workbook"]) - assert error is None, f"Unexpected error: {error}" - assert isinstance(parsed["workbook"], File) - assert parsed["workbook"].path == "/sandbox/output/workbook/result.xlsx" - - -class TestOutputDirParameter: - """Tests for output_dir parameter on PredictRLM.""" - - def test_output_dir_none_by_default(self): - from unittest.mock import MagicMock - - from predict_rlm import PredictRLM - - rlm = PredictRLM("query -> answer", sub_lm=MagicMock(), max_iterations=1) - assert rlm._output_dir is None - - def test_output_dir_stored_as_string(self): - from pathlib import Path - from unittest.mock import MagicMock - - from predict_rlm import PredictRLM - - rlm = PredictRLM( - "query -> answer", - sub_lm=MagicMock(), - max_iterations=1, - output_dir=Path("/tmp/test"), - ) - assert rlm._output_dir == "/tmp/test" - assert isinstance(rlm._output_dir, str) - - def test_output_dir_flows_to_file_plan(self): - class Sig(dspy.Signature): - query: str = dspy.InputField() - result: File = dspy.OutputField() - - from unittest.mock import MagicMock - - from predict_rlm import PredictRLM - - with tempfile.TemporaryDirectory() as tmpdir: - rlm = PredictRLM( - Sig, - sub_lm=MagicMock(), - max_iterations=1, - output_dir=tmpdir, - ) - plan, _ = rlm._prepare_file_io({"query": "test"}) - assert plan is not None - assert plan["write_dir"] == tmpdir - assert plan["output_field_map"]["result"]["host_dir"] == os.path.join( - tmpdir, "result" - ) - - -# -- Integration tests (require Deno + WASM sandbox) -- - - -@pytest.mark.integration -class TestFileIOIntegration: - """Full round-trip tests: mount file → sandbox reads it → sandbox writes output → sync back.""" - - def test_mount_and_read_file(self): - """Mount a host file and read its contents inside the sandbox.""" - from predict_rlm.backends import JspiBackend - - with tempfile.NamedTemporaryFile( - mode="w", suffix=".txt", delete=False - ) as f: - f.write("hello from host") - tmp_path = f.name - - try: - interpreter = JspiBackend( - preinstall_packages=False, - extra_read_paths=[tmp_path], - ) - try: - basename = os.path.basename(tmp_path) - interpreter._ensure_deno_process() - interpreter.mount_file_at(tmp_path, f"/sandbox/input/source/{basename}") - - result = interpreter.execute(f""" -content = open("/sandbox/input/source/{basename}").read() -print(content) -""") - assert "hello from host" in str(result) - finally: - interpreter.shutdown() - finally: - os.unlink(tmp_path) - - def test_write_and_sync_output_file(self): - """RLM writes a file in sandbox, sync it back to host.""" - from predict_rlm.backends import JspiBackend - - with tempfile.TemporaryDirectory() as tmpdir: - interpreter = JspiBackend( - preinstall_packages=False, - extra_write_paths=[tmpdir], - ) - try: - interpreter._ensure_deno_process() - interpreter.mkdir_p("/sandbox/output/result") - - interpreter.execute(""" -with open("/sandbox/output/result/output.txt", "w") as f: - f.write("generated by sandbox") -print("done") -""") - - # List files in output dir - files = interpreter.list_dir("/sandbox/output/result") - assert len(files) == 1 - assert files[0] == "/sandbox/output/result/output.txt" - - # Sync back to host - host_path = os.path.join(tmpdir, "output.txt") - interpreter.sync_file_to("/sandbox/output/result/output.txt", host_path) - - # Give sync_file a moment (it's a notification, not request-response) - import time - time.sleep(0.1) - - assert os.path.exists(host_path) - with open(host_path) as f: - assert f.read() == "generated by sandbox" - finally: - interpreter.shutdown() - - def test_full_roundtrip_mount_read_write_sync(self): - """Full round-trip: mount input → read in sandbox → transform → write output → sync.""" - from predict_rlm.backends import JspiBackend - - with tempfile.NamedTemporaryFile( - mode="w", suffix=".csv", delete=False - ) as f: - f.write("name,age\nAlice,30\nBob,25\n") - input_path = f.name - - with tempfile.TemporaryDirectory() as output_dir: - try: - interpreter = JspiBackend( - preinstall_packages=False, - extra_read_paths=[input_path], - extra_write_paths=[output_dir], - ) - try: - basename = os.path.basename(input_path) - interpreter._ensure_deno_process() - interpreter.mount_file_at( - input_path, f"/sandbox/input/data/{basename}" - ) - interpreter.mkdir_p("/sandbox/output/result") - - # Read CSV, transform, write output - interpreter.execute(f""" -import csv -import json - -with open("/sandbox/input/data/{basename}") as f: - reader = csv.DictReader(f) - rows = list(reader) - -# Transform: uppercase names -for row in rows: - row["name"] = row["name"].upper() - -with open("/sandbox/output/result/transformed.json", "w") as f: - json.dump(rows, f) - -print(f"Wrote {{len(rows)}} rows") -""") - - # List and sync - files = interpreter.list_dir("/sandbox/output/result") - assert len(files) == 1 - - host_output = os.path.join(output_dir, "transformed.json") - interpreter.sync_file_to(files[0], host_output) - - import time - time.sleep(0.1) - - # Verify - import json - - with open(host_output) as f: - data = json.load(f) - assert len(data) == 2 - assert data[0]["name"] == "ALICE" - assert data[1]["name"] == "BOB" - finally: - interpreter.shutdown() - finally: - os.unlink(input_path) - - def test_mkdir_p_creates_nested_dirs(self): - """mkdir_p creates deeply nested directories in MEMFS.""" - from predict_rlm.backends import JspiBackend - - interpreter = JspiBackend(preinstall_packages=False) - try: - interpreter._ensure_deno_process() - interpreter.mkdir_p("/sandbox/output/deep/nested/dir") - - result = interpreter.execute(""" -import os -exists = os.path.isdir("/sandbox/output/deep/nested/dir") -print(exists) -""") - assert "True" in str(result) - finally: - interpreter.shutdown() - - def test_list_dir_empty(self): - """list_dir on empty directory returns empty list.""" - from predict_rlm.backends import JspiBackend - - interpreter = JspiBackend(preinstall_packages=False) - try: - interpreter._ensure_deno_process() - interpreter.mkdir_p("/sandbox/output/empty") - files = interpreter.list_dir("/sandbox/output/empty") - assert files == [] - finally: - interpreter.shutdown() - - def test_list_dir_nonexistent(self): - """list_dir on nonexistent directory returns empty list.""" - from predict_rlm.backends import JspiBackend - - interpreter = JspiBackend(preinstall_packages=False) - try: - interpreter._ensure_deno_process() - files = interpreter.list_dir("/sandbox/output/doesnotexist") - assert files == [] - finally: - interpreter.shutdown() - - def test_mount_binary_file_roundtrip(self): - """Binary files survive the mount → read → write → sync round-trip.""" - from predict_rlm.backends import JspiBackend - - # Create a binary file with known bytes - binary_content = bytes(range(256)) - with tempfile.NamedTemporaryFile( - suffix=".bin", delete=False - ) as f: - f.write(binary_content) - input_path = f.name - - with tempfile.TemporaryDirectory() as output_dir: - try: - interpreter = JspiBackend( - preinstall_packages=False, - extra_read_paths=[input_path], - extra_write_paths=[output_dir], - ) - try: - basename = os.path.basename(input_path) - interpreter._ensure_deno_process() - interpreter.mount_file_at( - input_path, f"/sandbox/input/bin/{basename}" - ) - interpreter.mkdir_p("/sandbox/output/bin") - - # Copy binary file inside sandbox - interpreter.execute(f""" -data = open("/sandbox/input/bin/{basename}", "rb").read() -print(f"Read {{len(data)}} bytes") -with open("/sandbox/output/bin/copy.bin", "wb") as f: - f.write(data) -print("Written") -""") - - host_output = os.path.join(output_dir, "copy.bin") - interpreter.sync_file_to( - "/sandbox/output/bin/copy.bin", host_output - ) - - import time - time.sleep(0.1) - - with open(host_output, "rb") as f: - output_content = f.read() - assert output_content == binary_content - finally: - interpreter.shutdown() - finally: - os.unlink(input_path) + ] + ), + output_dir=output_dir, + max_iterations=1, + ) + + prediction = rlm(source=File.from_dir(str(source))) + + assert { + Path(item.path).relative_to(output_dir / "results").as_posix(): Path( + item.path + ).read_text(encoding="utf-8") + for item in prediction.results + } == {"first/result.txt": "FIRST", "second/result.txt": "SECOND"} diff --git a/tests/test_in_context.py b/tests/test_in_context.py index ed18dd05..852def31 100644 --- a/tests/test_in_context.py +++ b/tests/test_in_context.py @@ -1,23 +1,13 @@ -"""Tests for CtxStr prompt-injected string inputs.""" +"""Context-input precedence, composition boundaries, and run isolation.""" import asyncio -from types import MethodType from unittest.mock import MagicMock import dspy import pytest -import predict_rlm -import predict_rlm.in_context as in_context_module -from predict_rlm import CtxStr, CtxStrInputAdapter, PredictRLM, Skill -from predict_rlm.runtime import ( - BoundInput, - FieldDescriptor, - InputAdapter, - PreparedInput, - RunContext, - use_run_context, -) +from predict_rlm import CtxStr, CtxStrInputAdapter, PredictRLM +from predict_rlm.runtime import BoundInput, InputAdapter, PreparedInput class InContextSignature(dspy.Signature): @@ -49,21 +39,6 @@ async def _prepare_run( return ctx -def test_in_context_is_pydantic_string_schema(): - field = InContextSignature.input_fields["criteria"] - - assert field.annotation is CtxStr - assert InContextSignature.model_json_schema()["properties"]["criteria"]["type"] == "string" - - -def test_in_context_public_surface_is_only_ctx_str(): - assert in_context_module.__all__ == ["CtxStr", "CtxStrInputAdapter"] - assert "CtxStr" in predict_rlm.__all__ - assert "CtxStrInputAdapter" in predict_rlm.__all__ - assert "PromptContributor" not in predict_rlm.__all__ - assert "discover_runtime_modules" not in predict_rlm.__all__ - - @pytest.mark.asyncio async def test_ctx_str_resolves_to_builtin_adapter_ahead_of_generic_str_adapter(): rlm = PredictRLM( @@ -80,90 +55,6 @@ async def test_ctx_str_resolves_to_builtin_adapter_ahead_of_generic_str_adapter( assert ctx.input_bindings["query"].prepared.model_value == "prepared:QUESTION" -@pytest.mark.asyncio -async def test_ctx_str_subclass_resolves_to_builtin_ahead_of_generic_str_adapter(): - class SpecializedCtxStr(CtxStr): - pass - - class SpecializedSignature(dspy.Signature): - criteria: SpecializedCtxStr = dspy.InputField() - answer: str = dspy.OutputField() - - rlm = PredictRLM( - SpecializedSignature, - sub_lm=MagicMock(), - adapters=[PrefixStringInputAdapter()], - max_iterations=1, - ) - - ctx = await _prepare_run(rlm, {"criteria": "RULE"}) - - assert isinstance(ctx.input_bindings["criteria"].adapter, CtxStrInputAdapter) - assert "RULE" in str(ctx.state["generate_action"].signature.instructions) - - -@pytest.mark.asyncio -async def test_multiple_ctx_str_adapter_instances_append_one_ordered_section(): - class SpecializedCtxStr(CtxStr): - pass - - class MultipleCtxStrSignature(dspy.Signature): - first: CtxStr = dspy.InputField() - second: SpecializedCtxStr = dspy.InputField() - answer: str = dspy.OutputField() - - class SpecializedCtxStrAdapter(CtxStrInputAdapter): - name = "specialized_ctx_str" - value_type = SpecializedCtxStr - - rlm = PredictRLM( - MultipleCtxStrSignature, - sub_lm=MagicMock(), - adapters=[SpecializedCtxStrAdapter()], - max_iterations=1, - ) - - ctx = await _prepare_run(rlm, {"first": "FIRST", "second": "SECOND"}) - - for predictor_name in ("generate_action", "extract"): - prompt = str(ctx.state[predictor_name].signature.instructions) - assert prompt.count("## In-Context Inputs") == 1 - assert prompt.index("FIRST") < prompt.index("SECOND") - - -@pytest.mark.asyncio -async def test_default_ctx_str_prompt_excludes_custom_prompt_hook_fields(): - class SpecializedCtxStr(CtxStr): - pass - - class MultipleCtxStrSignature(dspy.Signature): - default: CtxStr = dspy.InputField() - custom: SpecializedCtxStr = dspy.InputField() - answer: str = dspy.OutputField() - - class CustomPromptAdapter(CtxStrInputAdapter): - name = "custom_prompt_ctx_str" - value_type = SpecializedCtxStr - - def append_prompt(self, prompt, field, prepared, ctx): - return f"{prompt}\n\nCUSTOM:{prepared.model_value}" - - rlm = PredictRLM( - MultipleCtxStrSignature, - sub_lm=MagicMock(), - adapters=[CustomPromptAdapter()], - max_iterations=1, - ) - - ctx = await _prepare_run(rlm, {"default": "DEFAULT", "custom": "CUSTOM-VALUE"}) - - for predictor_name in ("generate_action", "extract"): - prompt = str(ctx.state[predictor_name].signature.instructions) - assert prompt.count("## In-Context Inputs") == 1 - assert prompt.count("CUSTOM-VALUE") == 1 - assert "CUSTOM:CUSTOM-VALUE" in prompt - - def test_independent_same_name_ctx_str_adapter_is_rejected_at_construction(): class UnsafeCtxStrAdapter(InputAdapter[CtxStr]): name = "ctx_str" @@ -181,26 +72,6 @@ async def prepare(self, field, value, ctx): ) -@pytest.mark.asyncio -async def test_same_name_ctx_str_adapter_replaces_builtin_and_owns_prompt(): - class ReplacementCtxStrInputAdapter(CtxStrInputAdapter): - async def prepare(self, field, value, ctx): - return PreparedInput(model_value=f"replacement:{value}") - - rlm = PredictRLM( - InContextSignature, - sub_lm=MagicMock(), - adapters=[ReplacementCtxStrInputAdapter()], - max_iterations=1, - ) - - ctx = await _prepare_run(rlm, {"criteria": "RULE", "query": "QUESTION"}) - prompt = str(ctx.state["generate_action"].signature.instructions) - - assert type(ctx.input_bindings["criteria"].adapter) is ReplacementCtxStrInputAdapter - assert "replacement:RULE" in prompt - - @pytest.mark.asyncio async def test_distinct_exact_ctx_str_adapter_conflicts_with_builtin(): class OtherCtxStrAdapter(InputAdapter[CtxStr]): @@ -221,22 +92,6 @@ async def prepare(self, field, value, ctx): await _prepare_run(rlm, {"criteria": "RULE", "query": "QUESTION"}) -@pytest.mark.asyncio -async def test_kernel_does_not_build_ctx_str_instructions(monkeypatch): - monkeypatch.setattr( - in_context_module, - "_build_in_context_instructions", - MagicMock(side_effect=AssertionError("kernel special case called")), - raising=False, - ) - rlm = PredictRLM(InContextSignature, sub_lm=MagicMock(), max_iterations=1) - - ctx = await _prepare_run(rlm, {"criteria": "RULE", "query": "QUESTION"}) - - assert "RULE" in str(ctx.state["generate_action"].signature.instructions) - - - @pytest.mark.asyncio async def test_in_context_delimiters_avoid_prepared_value_collisions(): nominal_closing_marker = '' @@ -260,96 +115,6 @@ async def test_in_context_delimiters_avoid_prepared_value_collisions(): assert not instructions.rstrip().endswith(nominal_closing_marker) -@pytest.mark.asyncio -async def test_runtime_in_context_input_uses_run_local_predictors_and_repl_value(): - rlm = PredictRLM(InContextSignature, sub_lm=MagicMock(), max_iterations=1) - original_action = rlm.generate_action - original_extract = rlm.extract - criteria = "Always mention the controlling rule." - - ctx = await _prepare_run( - rlm, - {"criteria": criteria, "query": "What matters?"}, - ) - - action = str(ctx.state["generate_action"].signature.instructions) - extract = str(ctx.state["extract"].signature.instructions) - assert ctx.input_bindings["criteria"].prepared.model_value == criteria - assert action.count(criteria) == 1 - assert extract.count(criteria) == 1 - assert action.rstrip().endswith('') - assert extract.rstrip().endswith('') - assert rlm.generate_action is original_action - assert rlm.extract is original_extract - - -@pytest.mark.asyncio -async def test_runtime_in_context_instructions_follow_files_and_skills(): - skill = Skill(name="domain", instructions="Skill block") - rlm = PredictRLM( - InContextSignature, - sub_lm=MagicMock(), - max_iterations=1, - skills=[skill], - ) - - ctx = await _prepare_run( - rlm, - {"criteria": "Criteria block", "query": "What matters?"}, - file_instructions="File block", - ) - action = str(ctx.state["generate_action"].signature.instructions) - extract = str(ctx.state["extract"].signature.instructions) - - assert action.index("File block") < action.index("Skill block") - assert action.index("Skill block") < action.index("Criteria block") - assert extract.index("File block") < extract.index("Skill block") - assert extract.index("Skill block") < extract.index("Criteria block") - - -@pytest.mark.asyncio -async def test_in_context_preserves_invocation_signature_builder_override(): - class TransformAdapter(CtxStrInputAdapter): - name = "ctx_str" - value_type = CtxStr - - async def prepare(self, field, value, ctx): - return PreparedInput(model_value=value) - - def _transform_prompt_signature(self, signature, field, prepared, ctx): - signature.instructions = f"transformed:{prepared.model_value}" - return signature - - class TrackingPredictRLM(PredictRLM): - def _build_signatures_with_files(self, file_instructions): - self.file_builder_calls = getattr(self, "file_builder_calls", 0) + 1 - action, extract = super()._build_signatures_with_files(file_instructions) - action.builder_marker = "action" - extract.builder_marker = "extract" - return action, extract - - rlm = TrackingPredictRLM( - InContextSignature, - sub_lm=MagicMock(), - adapters=[TransformAdapter()], - max_iterations=1, - ) - - ctx = await _prepare_run( - rlm, - {"criteria": "Criteria block", "query": "What matters?"}, - file_instructions="File block", - ) - - assert rlm.file_builder_calls == 1 - assert ctx.state["generate_action"].builder_marker == "action" - assert ctx.state["extract"].builder_marker == "extract" - for predictor_name in ("generate_action", "extract"): - assert "transformed:Criteria block" in str( - ctx.state[predictor_name].signature.instructions - ) - - @pytest.mark.asyncio async def test_concurrent_in_context_runs_do_not_cross_contaminate_predictors(): rlm = PredictRLM(InContextSignature, sub_lm=MagicMock(), max_iterations=1) @@ -376,30 +141,9 @@ async def test_in_context_rejects_non_string_runtime_value(): await _prepare_run(rlm, {"criteria": 123, "query": "What matters?"}) -@pytest.mark.asyncio -async def test_generic_string_input_adapter_does_not_own_ctx_str_value(): - rlm = PredictRLM( - InContextSignature, - sub_lm=MagicMock(), - adapters=[PrefixStringInputAdapter()], - max_iterations=1, - ) - - ctx = await _prepare_run( - rlm, - {"criteria": "RAW-RULE", "query": "What matters?"}, - ) - action = str(ctx.state["generate_action"].signature.instructions) - - assert ctx.input_bindings["criteria"].prepared.model_value == "RAW-RULE" - assert "\nRAW-RULE\n" in action - assert "prepared:RAW-RULE" not in action - - @pytest.mark.asyncio async def test_ctx_str_prompt_uses_final_bound_custom_adapter_value(): class BindingAdapter(CtxStrInputAdapter): - async def prepare(self, field, value, ctx): return PreparedInput(model_value=f"prepared:{value}") @@ -424,73 +168,6 @@ async def bind(self, field, prepared, ctx, session): assert "\nprepared:RULE\n" not in action -@pytest.mark.asyncio -async def test_input_adapter_append_prompt_sees_current_prompt_for_action_and_extract(): - seen = [] - - class PromptAdapter(PrefixStringInputAdapter): - def append_prompt(self, prompt, field, prepared, ctx): - seen.append((prompt, field.name, prepared.model_value, ctx.run_id)) - return f"{prompt}\n\nHOOK:{field.name}:{prepared.model_value}" - - rlm = PredictRLM( - InContextSignature, - sub_lm=MagicMock(), - adapters=[PromptAdapter()], - max_iterations=1, - ) - ctx = await _prepare_run( - rlm, - {"criteria": "RULE", "query": "QUESTION"}, - ) - - action = str(ctx.state["generate_action"].signature.instructions) - extract = str(ctx.state["extract"].signature.instructions) - assert len(seen) == 2 - assert [(field, value) for _, field, value, _ in seen] == [ - ("query", "prepared:QUESTION"), - ("query", "prepared:QUESTION"), - ] - assert all(run_id == ctx.run_id for *_, run_id in seen) - assert "HOOK:query:prepared:QUESTION" in action - assert "HOOK:query:prepared:QUESTION" in extract - - -@pytest.mark.asyncio -async def test_prompt_hooks_chain_in_signature_field_order(): - class CriteriaAdapter(CtxStrInputAdapter): - - async def prepare(self, field, value, ctx): - return PreparedInput(model_value=value) - - def append_prompt(self, prompt, field, prepared, ctx): - return f"{prompt}\nFIRST" - - class QueryAdapter(InputAdapter[str]): - name = "value" - value_type = str - fallback = True - - async def prepare(self, field, value, ctx): - return PreparedInput(model_value=value) - - def append_prompt(self, prompt, field, prepared, ctx): - assert "FIRST" in prompt - return f"{prompt}\nSECOND" - - rlm = PredictRLM( - InContextSignature, - sub_lm=MagicMock(), - adapters=[QueryAdapter(), CriteriaAdapter()], - max_iterations=1, - ) - ctx = await _prepare_run(rlm, {"criteria": "RULE", "query": "QUESTION"}) - - for key in ("generate_action", "extract"): - prompt = str(ctx.state[key].signature.instructions) - assert prompt.index("FIRST") < prompt.index("SECOND") - - @pytest.mark.asyncio async def test_in_place_prompt_signature_transform_is_run_local(): transformed_signatures = {} @@ -547,29 +224,6 @@ def _transform_prompt_signature(self, signature, field, prepared, ctx): assert rlm.signature.input_fields["query"].json_schema_extra["desc"] == original_description -@pytest.mark.asyncio -async def test_instance_bound_prompt_signature_transform_is_applied(): - adapter = PrefixStringInputAdapter() - - def transform_prompt_signature(self, signature, field, prepared, ctx): - signature.instructions += f"\ninstance:{field.name}:{prepared.model_value}" - return signature - - adapter._transform_prompt_signature = MethodType(transform_prompt_signature, adapter) - rlm = PredictRLM( - InContextSignature, - sub_lm=MagicMock(), - adapters=[adapter], - max_iterations=1, - ) - - ctx = await _prepare_run(rlm, {"criteria": "RULE", "query": "QUESTION"}) - - for predictor_name in ("generate_action", "extract"): - prompt = str(ctx.state[predictor_name].signature.instructions) - assert "instance:query:prepared:QUESTION" in prompt - - @pytest.mark.asyncio @pytest.mark.parametrize("hook", ["append_prompt", "_transform_prompt_signature"]) async def test_invalid_prompt_hook_return_type_fails_clearly(hook): @@ -588,146 +242,6 @@ class InvalidAdapter(PrefixStringInputAdapter): await _prepare_run(rlm, {"criteria": "RULE", "query": "QUESTION"}) -def test_removed_prompt_contributor_magic_is_absent(): - import predict_rlm.runtime as runtime_module - - assert not hasattr(runtime_module, "_PromptContributor") - assert not hasattr(runtime_module, "_discover_annotation_prompt_contributors") - assert not hasattr(CtxStr, "_predict_rlm_prompt_contributor") - - -def test_input_adapter_prompt_hook_defaults_are_noops(): - adapter = PrefixStringInputAdapter() - prepared = PreparedInput(model_value="value") - field = FieldDescriptor("query", str) - ctx = RunContext(MagicMock(), {}) - - assert adapter.append_prompt("prompt", field, prepared, ctx) == "prompt" - assert adapter._transform_prompt_signature(InContextSignature, field, prepared, ctx) is InContextSignature - - -@pytest.mark.asyncio -async def test_default_prompt_hook_preserves_custom_action_predictor(): - class PlainSignature(dspy.Signature): - query: str = dspy.InputField() - answer: str = dspy.OutputField() - - class CustomAction: - pass - - action = CustomAction() - rlm = PredictRLM( - PlainSignature, - sub_lm=MagicMock(), - adapters=[PrefixStringInputAdapter()], - max_iterations=1, - ) - rlm.generate_action = action - - ctx = await _prepare_run(rlm, {"query": "QUESTION"}) - - assert ctx.state["generate_action"] is action - - -@pytest.mark.asyncio -async def test_delegating_sync_iteration_uses_run_local_in_context_predictor(): - class DelegatingPredictRLM(PredictRLM): - def _execute_iteration(self, *args, **kwargs): - return super()._execute_iteration(*args, **kwargs) - - rlm = DelegatingPredictRLM( - InContextSignature, - sub_lm=MagicMock(), - max_iterations=1, - ) - ctx = await _prepare_run( - rlm, - {"criteria": "RUN-RULE", "query": "What matters?"}, - ) - run_action = MagicMock( - return_value=dspy.Prediction(reasoning="done", code="SUBMIT(answer='done')") - ) - ctx.state["generate_action"] = run_action - rlm.generate_action = MagicMock(side_effect=AssertionError("used shared predictor")) - rlm._record_action_generation_ok = MagicMock(return_value=None) - rlm._prepare_iteration_execution = MagicMock(return_value=("code", False)) - rlm._execute_iteration_code = MagicMock(return_value="result") - rlm._complete_iteration_execution = MagicMock(return_value="done") - - async with use_run_context(ctx): - result = await rlm._aexecute_iteration(None, [], [], 0, {}, ["answer"]) - - assert result == "done" - run_action.assert_called_once() - rlm.generate_action.assert_not_called() - - -@pytest.mark.asyncio -async def test_delegating_extract_fallback_uses_run_local_in_context_predictor(): - class DelegatingPredictRLM(PredictRLM): - def _extract_fallback(self, *args, **kwargs): - return super()._extract_fallback(*args, **kwargs) - - rlm = DelegatingPredictRLM( - InContextSignature, - sub_lm=MagicMock(), - max_iterations=1, - ) - ctx = await _prepare_run( - rlm, - {"criteria": "RUN-RULE", "query": "What matters?"}, - ) - run_extract = MagicMock(return_value=dspy.Prediction(answer="run-local")) - ctx.state["extract"] = run_extract - rlm.extract = MagicMock(side_effect=AssertionError("used shared predictor")) - - async with use_run_context(ctx): - result = await rlm._aextract_fallback_for_run([], [], ["answer"]) - - assert result.answer == "run-local" - run_extract.assert_called_once() - rlm.extract.assert_not_called() - - -@pytest.mark.asyncio -async def test_sync_extract_fallback_calls_cooperative_later_mro_with_run_local_copy(): - fallback_instances = [] - - class CooperativeFallbackRLM(dspy.RLM): - def _extract_fallback(self, variables, history, output_field_names): - fallback_instances.append(self) - return super()._extract_fallback(variables, history, output_field_names) - - class CooperativePredictRLM(PredictRLM, CooperativeFallbackRLM): - pass - - rlm = CooperativePredictRLM( - InContextSignature, - sub_lm=MagicMock(), - max_iterations=1, - ) - ctx = await _prepare_run( - rlm, - {"criteria": "RUN-RULE", "query": "What matters?"}, - ) - run_extract = MagicMock(return_value=dspy.Prediction(answer="run-local")) - ctx.state["extract"] = run_extract - rlm.extract = MagicMock(side_effect=AssertionError("used shared predictor")) - - async with use_run_context(ctx): - result = rlm._extract_fallback_for_run([], [], ["answer"]) - - assert result.answer == "run-local" - assert result.trajectory == [] - assert result.final_reasoning == "Extract forced final output" - assert len(fallback_instances) == 1 - assert type(fallback_instances[0]) is CooperativePredictRLM - assert fallback_instances[0] is not rlm - assert fallback_instances[0].extract is run_extract - run_extract.assert_called_once_with(variables_info=[], repl_history=[]) - rlm.extract.assert_not_called() - - def test_in_context_is_input_only(): class BadOutput(dspy.Signature): prompt: str = dspy.InputField() @@ -744,12 +258,3 @@ class BadOptionalInput(dspy.Signature): with pytest.raises(TypeError, match="annotated directly"): PredictRLM(BadOptionalInput, sub_lm=MagicMock(), max_iterations=1) - - -def test_in_context_rejects_list_input_annotation(): - class BadListInput(dspy.Signature): - criteria: list[CtxStr] = dspy.InputField() - answer: str = dspy.OutputField() - - with pytest.raises(TypeError, match="annotated directly"): - PredictRLM(BadListInput, sub_lm=MagicMock(), max_iterations=1) diff --git a/tests/test_interpreter.py b/tests/test_interpreter.py index cd6562da..993ceea2 100644 --- a/tests/test_interpreter.py +++ b/tests/test_interpreter.py @@ -1,12 +1,22 @@ """Tests for JspiBackend with concurrent async tool execution.""" +import asyncio +import shutil +from pathlib import Path +from unittest.mock import AsyncMock, MagicMock + +import dspy import pytest from dspy.primitives.code_interpreter import CodeInterpreterError, FinalOutput +from predict_rlm import PredictRLM, Workspace from predict_rlm.backends import JspiBackend from predict_rlm.backends.base import SandboxFatalError -pytestmark = pytest.mark.integration +pytestmark = [ + pytest.mark.integration, + pytest.mark.skipif(shutil.which("deno") is None, reason="JSPI tests require Deno"), +] class TestSubmitDefaults: @@ -43,32 +53,6 @@ def test_bare_submit_without_default_still_errors(self): class TestCodeFenceStripping: - """Tests for code fence extraction (```repl preferred, ```python fallback).""" - - def test_strip_repl_fence(self): - """Code wrapped in ```repl fence is extracted and executed.""" - interpreter = JspiBackend(preinstall_packages=False) - try: - result = interpreter.execute("""```repl -x = 2 + 2 -print(x) -```""") - assert "4" in str(result) - finally: - interpreter.shutdown() - - def test_strip_repl_fence_with_text_before(self): - """```repl fence with explanatory text before is handled.""" - interpreter = JspiBackend(preinstall_packages=False) - try: - result = interpreter.execute("""Here's the code to run: -```repl -print("extracted") -```""") - assert "extracted" in str(result) - finally: - interpreter.shutdown() - def test_repl_fence_handles_inline_backticks(self): """Inline ``` (not on own line) inside code is preserved.""" interpreter = JspiBackend(preinstall_packages=False) @@ -82,76 +66,6 @@ def test_repl_fence_handles_inline_backticks(self): finally: interpreter.shutdown() - def test_fallback_python_fence(self): - """Falls back to ```python fence for backwards compatibility.""" - interpreter = JspiBackend(preinstall_packages=False) - try: - result = interpreter.execute("""```python -print("fallback") -```""") - assert "fallback" in str(result) - finally: - interpreter.shutdown() - - def test_fallback_bare_fence(self): - """Falls back to bare ``` fence for backwards compatibility.""" - interpreter = JspiBackend(preinstall_packages=False) - try: - result = interpreter.execute("""``` -print(42) -```""") - assert "42" in str(result) - finally: - interpreter.shutdown() - - def test_no_fence_unchanged(self): - """Code without fences executes normally.""" - interpreter = JspiBackend(preinstall_packages=False) - try: - result = interpreter.execute("print('no fence')") - assert "no fence" in str(result) - finally: - interpreter.shutdown() - - def test_prints_before_runtime_error_are_in_error(self): - interpreter = JspiBackend(preinstall_packages=False) - try: - with pytest.raises(CodeInterpreterError) as exc_info: - interpreter.execute("print('before failure')\nraise ValueError('bad')") - finally: - interpreter.shutdown() - - assert "before failure" in str(exc_info.value) - assert "ValueError" in str(exc_info.value) - assert getattr(exc_info.value, "partial_output") == "before failure\n" - - def test_verbose_prints_partial_output_before_error(self, capsys): - interpreter = JspiBackend(preinstall_packages=False, verbose=True) - try: - with pytest.raises(CodeInterpreterError): - interpreter.execute("print('before failure')\nraise ValueError('bad')") - finally: - interpreter.shutdown() - - stderr = capsys.readouterr().err - assert "output:" in stderr - assert "before failure" in stderr - assert "error (ValueError):" in stderr - assert "bad" in stderr - - def test_double_fence_handled(self): - """Double fences (model outputs ```...```\\n```) are handled correctly.""" - interpreter = JspiBackend(preinstall_packages=False) - try: - # Model sometimes outputs a trailing bare ``` after the closing fence - result = interpreter.execute("""```repl -print("double fence") -``` -```""") - assert "double fence" in str(result) - finally: - interpreter.shutdown() - def test_multiple_repl_blocks(self): """Multiple ```repl blocks are all extracted and executed in order.""" interpreter = JspiBackend(preinstall_packages=False) @@ -180,2767 +94,244 @@ def test_multiple_repl_blocks(self): finally: interpreter.shutdown() - def test_multiple_repl_blocks_with_async(self): - """Multiple ```repl blocks with async code work correctly.""" - import asyncio - - call_log = [] - - async def log_tool(msg): - call_log.append(msg) - await asyncio.sleep(0.01) - return f"logged: {msg}" - - interpreter = JspiBackend( - preinstall_packages=False, - tools={"log_tool": log_tool}, - ) - try: - result = interpreter.execute("""First block: -```repl -result1 = await log_tool("block1") -print(result1) -``` - -Second block: -```repl -result2 = await log_tool("block2") -print(result2) -``` - -Third block uses results from previous: -```repl -print(f"Results: {result1}, {result2}") -```""") - assert "logged: block1" in str(result) - assert "logged: block2" in str(result) - assert "Results: logged: block1, logged: block2" in str(result) - assert call_log == ["block1", "block2"] - finally: - interpreter.shutdown() - - -class TestJspiBackend: - """Tests that JspiBackend executes code correctly.""" - - def test_interpreter_executes_python_code(self): - """JspiBackend executes Python code and returns output.""" - interpreter = JspiBackend(preinstall_packages=False) - try: - result = interpreter.execute("x = 2 + 2\nprint(x)") - assert "4" in str(result) - finally: - interpreter.shutdown() - - def test_interpreter_state_persists(self): - """Variables persist between code executions in the same interpreter.""" - interpreter = JspiBackend(preinstall_packages=False) - try: - interpreter.execute("my_list = [1, 2, 3]") - interpreter.execute("my_list.append(4)") - result = interpreter.execute("print(len(my_list))") - assert "4" in str(result) - finally: - interpreter.shutdown() - - def test_interpreter_can_use_stdlib(self): - """Interpreter has access to standard library modules.""" - interpreter = JspiBackend(preinstall_packages=False) - try: - result = interpreter.execute(""" -import re -import json -import math - -data = {"value": math.pi} -text = json.dumps(data) -match = re.search(r'"value": ([0-9.]+)', text) -print(match.group(1)[:4]) -""") - assert "3.14" in str(result) - finally: - interpreter.shutdown() - - def test_interpreter_has_jspi_flag_when_needed(self): - """JspiBackend includes the JSPI V8 flag only when V8 < 13.7.""" - from predict_rlm.backends.jspi.backend import _needs_jspi_flag - - interpreter = JspiBackend(preinstall_packages=False) - try: - if _needs_jspi_flag(): - assert "--v8-flags=--experimental-wasm-jspi" in interpreter.deno_command - else: - assert "--v8-flags=--experimental-wasm-jspi" not in interpreter.deno_command - finally: - interpreter.shutdown() - - def test_interpreter_has_pypi_network_by_default(self): - """JspiBackend has PyPI network access by default for package installation.""" - interpreter = JspiBackend() - try: - net_flags = [arg for arg in interpreter.deno_command if "--allow-net" in arg] - assert len(net_flags) == 1 - # PyPI domains are required for micropip to install packages - assert "pypi.org" in net_flags[0] - assert "files.pythonhosted.org" in net_flags[0] - finally: - interpreter.shutdown() - - def test_interpreter_with_allowed_domains(self): - """JspiBackend can be configured with allowed domains.""" - interpreter = JspiBackend( - preinstall_packages=False, allowed_domains=["api.example.com", "cdn.example.com"] - ) - try: - net_flags = [arg for arg in interpreter.deno_command if "--allow-net" in arg] - assert len(net_flags) == 1 - assert "api.example.com" in net_flags[0] - assert "cdn.example.com" in net_flags[0] - finally: - interpreter.shutdown() - - -class TestInterpreterWithTools: - """Tests that interpreter can call registered tools.""" - - def test_interpreter_calls_tool_from_code(self): - """Interpreter can call a tool from within executed code.""" - call_log: list[dict] = [] - - def predict(signature: str, **kwargs) -> dict: - call_log.append({"signature": signature, "kwargs": kwargs}) - outputs = signature.split("->")[1].strip().split(",") - return {out.strip(): f"Result for {out.strip()}" for out in outputs} - - interpreter = JspiBackend(preinstall_packages=False, tools={"predict": predict}) - try: - output = interpreter.execute(""" -result = await predict("question -> answer", question="What is 2+2?") -print(result["answer"]) -""") - assert "Result for answer" in str(output) - assert len(call_log) == 1 - assert call_log[0]["kwargs"]["question"] == "What is 2+2?" - finally: - interpreter.shutdown() - - def test_interpreter_can_process_tool_output(self): - """Interpreter can process tool output with Python code.""" - - def get_prices() -> str: - return "Item 1: $10.00\nItem 2: $20.00\nItem 3: $15.00" - - interpreter = JspiBackend( - preinstall_packages=False, tools={"get_prices": get_prices} - ) - try: - output = interpreter.execute(""" -import re -text = await get_prices() -prices = re.findall(r'\\$(\\d+\\.\\d+)', text) -total = sum(float(p) for p in prices) -print(f"Total: ${total:.2f}") -""") - assert "Total: $45.00" in str(output) - finally: - interpreter.shutdown() +class TestPydanticSerialization: + def test_pydantic_model_tool_result_is_mapping(self): + from pydantic import BaseModel - def test_multiple_tools(self): - """Multiple tools can be registered and called.""" + class Source(BaseModel): + title: str - def fetch_data() -> str: - return "ABC123" + class Retrieval(BaseModel): + model_name: str + sources: list[Source] - def format_id(raw_id: str) -> str: - return f"ID-{raw_id.strip()}" + def retrieve() -> Retrieval: + return Retrieval( + model_name="test-index", + sources=[Source(title="PredictRLM documentation")], + ) - interpreter = JspiBackend( - preinstall_packages=False, - tools={ - "fetch_data": fetch_data, - "format_id": format_id, - }, - ) + interpreter = JspiBackend(tools={"retrieve": retrieve}) try: output = interpreter.execute(""" -raw = await fetch_data() -formatted = await format_id(raw) -print(formatted) +result = await retrieve() +print(result["model_name"]) +print(result["sources"][0]["title"]) """) - assert "ID-ABC123" in str(output) + assert "test-index" in str(output) + assert "PredictRLM documentation" in str(output) finally: interpreter.shutdown() + def test_nested_pydantic_models(self): + """Nested Pydantic models are serialized correctly.""" + received_data = [] -class TestConcurrentToolExecution: - """Tests for concurrent/parallel tool execution using async.""" - - def test_async_tool_calls_run_concurrently(self): - """Async tool calls via asyncio.gather() run in parallel.""" - import time - - def slow_tool(item_id: str) -> str: - time.sleep(0.005) - return f"Result for {item_id}" + def process_nested(data: dict) -> str: + received_data.append(data) + return f"Got {len(data.get('items', []))} items" - interpreter = JspiBackend(preinstall_packages=False, tools={"slow_tool": slow_tool}) + interpreter = JspiBackend(tools={"process_nested": process_nested}) try: output = interpreter.execute(""" -import asyncio +from pydantic import BaseModel +from typing import List, Optional -tasks = [slow_tool(f"item_{i}") for i in range(3)] -results = await asyncio.gather(*tasks) -print(results) -""") - assert "Result for item_0" in str(output) - assert "Result for item_1" in str(output) - assert "Result for item_2" in str(output) - finally: - interpreter.shutdown() +class Item(BaseModel): + id: int + name: str + price: float - def test_single_tool_call(self): - """Single tool call with await works correctly.""" +class Order(BaseModel): + order_id: str + customer: str + items: List[Item] + notes: Optional[str] = None - def my_tool(value: str) -> str: - return f"Got: {value}" +order = Order( + order_id="ORD-123", + customer="Bob", + items=[ + Item(id=1, name="Widget", price=9.99), + Item(id=2, name="Gadget", price=19.99), + ], + notes="Rush delivery" +) - interpreter = JspiBackend(preinstall_packages=False, tools={"my_tool": my_tool}) - try: - output = interpreter.execute(""" -result = await my_tool("test_value") +result = await process_nested(order) print(result) """) - assert "Got: test_value" in str(output) - finally: - interpreter.shutdown() - - def test_sequential_and_parallel_tools(self): - """Can use sequential await and parallel asyncio.gather in same code.""" - call_log: list[str] = [] - - def tracker(msg: str) -> str: - call_log.append(msg) - return f"Tracked: {msg}" - - interpreter = JspiBackend(preinstall_packages=False, tools={"tracker": tracker}) - try: - output = interpreter.execute(""" -import asyncio - -# Sequential call -first_result = await tracker("first") -print(f"First: {first_result}") - -# Parallel calls with asyncio.gather -parallel_results = await asyncio.gather( - tracker("parallel_1"), - tracker("parallel_2"), -) -print(f"Parallel: {parallel_results}") -""") - assert "Tracked: first" in str(output) - assert "Tracked: parallel_1" in str(output) - assert "Tracked: parallel_2" in str(output) - assert len(call_log) == 3 + assert "Got 2 items" in str(output) + assert len(received_data) == 1 + assert received_data[0]["order_id"] == "ORD-123" + assert len(received_data[0]["items"]) == 2 + assert received_data[0]["items"][0]["name"] == "Widget" finally: interpreter.shutdown() + def test_non_serializable_falls_back_to_string(self): + """Non-serializable objects are converted to strings gracefully.""" + received_args = [] -class TestPreinstalledPackages: - """Tests that pandas and pydantic are pre-installed in the sandbox.""" - - def test_pandas_is_available(self): - """Pandas can be imported and used in the sandbox.""" - interpreter = JspiBackend() - try: - output = interpreter.execute(""" -import pandas as pd - -df = pd.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]}) -print(f"Shape: {df.shape}") -print(f"Sum of a: {df['a'].sum()}") -""") - assert "Shape: (3, 2)" in str(output) - assert "Sum of a: 6" in str(output) - finally: - interpreter.shutdown() + def receive_anything(x) -> str: + received_args.append(x) + return f"Got type: {type(x).__name__}" - def test_pydantic_is_available(self): - """Pydantic can be imported and used for validation in the sandbox.""" - interpreter = JspiBackend() + interpreter = JspiBackend(tools={"receive_anything": receive_anything}) try: + # Pass a coroutine (non-standard object) - should be converted to string output = interpreter.execute(""" -from pydantic import BaseModel, Field - -class Item(BaseModel): - name: str - price: float = Field(gt=0) +async def dummy(): + return 42 -item = Item(name="Widget", price=9.99) -print(f"Item: {item.name} @ ${item.price}") -print(f"Model fields: {list(item.model_fields.keys())}") +coro = dummy() +result = await receive_anything(coro) +print(result) +# Clean up +coro.close() """) - assert "Item: Widget @ $9.99" in str(output) - assert "name" in str(output) - assert "price" in str(output) - finally: - interpreter.shutdown() - - def test_pydantic_validation_works(self): - """Pydantic validation errors are raised correctly in the sandbox.""" - interpreter = JspiBackend() - try: - output = interpreter.execute(""" -from pydantic import BaseModel, ValidationError - -class User(BaseModel): - name: str - age: int + # Should succeed with string representation + assert "Got type: str" in str(output) + assert len(received_args) == 1 + # The coroutine should have been converted to its string repr + assert "coroutine" in str(received_args[0]).lower() -try: - user = User(name="Alice", age="not_a_number") -except ValidationError as e: - print(f"Validation error count: {e.error_count()}") - print("Validation works!") + # Tool should still work normally + output2 = interpreter.execute(""" +result = await receive_anything({"normal": "dict"}) +print(result) """) - assert "Validation works!" in str(output) + assert "Got type: dict" in str(output2) finally: interpreter.shutdown() -class TestToolPersistence: - """Tests that tools persist across multiple executions and heavy workloads.""" - - def test_tools_persist_across_executions(self): - """Tools remain available after multiple code executions.""" - - def my_tool(value: str) -> str: - return f"Got: {value}" +class TestNoneValueSerialization: + def test_pydantic_model_with_none_fields_injected_as_variable(self): + """Pydantic models with None fields are accessible in the sandbox.""" + from pydantic import BaseModel, Field - interpreter = JspiBackend(preinstall_packages=False, tools={"my_tool": my_tool}) - try: - # First execution - use the tool - output1 = interpreter.execute(""" -result = await my_tool("first") -print(result) -""") - assert "Got: first" in str(output1) - - # Second execution - define some variables (potentially corrupting state) - interpreter.execute(""" -import asyncio -data = [i * 2 for i in range(100)] -total = sum(data) -""") - - # Third execution - tool should still work - output3 = interpreter.execute(""" -result = await my_tool("third") -print(result) -""") - assert "Got: third" in str(output3) - finally: - interpreter.shutdown() - - def test_tools_survive_heavy_async_workload(self): - """Tools remain available after heavy concurrent async operations.""" - import time - - call_count = 0 - - def counter_tool(item_id: str) -> str: - nonlocal call_count - call_count += 1 - time.sleep(0.001) # Small delay to simulate work - return f"Processed: {item_id}" - - interpreter = JspiBackend( - preinstall_packages=False, tools={"counter_tool": counter_tool} - ) - try: - # Heavy concurrent workload - output1 = interpreter.execute(""" -import asyncio - -# Run many parallel calls -tasks = [counter_tool(f"item_{i}") for i in range(20)] -results = await asyncio.gather(*tasks) -print(f"Completed {len(results)} calls") -""") - assert "Completed 20 calls" in str(output1) - assert call_count == 20 - - # After heavy workload, tool should still be available - output2 = interpreter.execute(""" -# Check tool is still callable -result = await counter_tool("after_workload") -print(result) -""") - assert "Processed: after_workload" in str(output2) - finally: - interpreter.shutdown() - - def test_tools_in_persistent_module(self): - """Tools are stored in _repl_tools module and can be recovered.""" - - def check_tool() -> str: - return "check_ok" - - interpreter = JspiBackend( - preinstall_packages=False, tools={"check_tool": check_tool} - ) - try: - # Verify tool exists in module - output = interpreter.execute(""" -import sys - -# Check module exists -has_module = '_repl_tools' in sys.modules -print(f"Module exists: {has_module}") - -if has_module: - _repl_tools = sys.modules['_repl_tools'] - has_tool = hasattr(_repl_tools, 'check_tool') - print(f"Tool in module: {has_tool}") - -# Call tool via global -result = await check_tool() -print(f"Result: {result}") -""") - assert "Module exists: True" in str(output) - assert "Tool in module: True" in str(output) - assert "Result: check_ok" in str(output) - finally: - interpreter.shutdown() - - -class TestToolFailures: - """Tests for tool error handling and recovery.""" - - def test_tool_exception_is_catchable(self): - """Exceptions from tools can be caught in Python code.""" - - def failing_tool() -> str: - raise ValueError("Tool failed intentionally") - - interpreter = JspiBackend( - preinstall_packages=False, tools={"failing_tool": failing_tool} - ) - try: - output = interpreter.execute(""" -try: - result = await failing_tool() - print("Should not reach here") -except Exception as e: - print(f"Caught exception: {type(e).__name__}") - print(f"Message contains 'failed': {'failed' in str(e)}") -""") - assert "Caught exception:" in str(output) - assert "Message contains 'failed': True" in str(output) - finally: - interpreter.shutdown() - - def test_tools_work_after_exception(self): - """Tools continue working after one raises an exception.""" - call_count = 0 - - def maybe_fail(should_fail: bool) -> str: - nonlocal call_count - call_count += 1 - if should_fail: - raise RuntimeError("Intentional failure") - return "success" - - interpreter = JspiBackend( - preinstall_packages=False, tools={"maybe_fail": maybe_fail} - ) - try: - # First call succeeds - output1 = interpreter.execute(""" -result = await maybe_fail(False) -print(f"First call: {result}") -""") - assert "First call: success" in str(output1) - - # Second call fails (caught) - output2 = interpreter.execute(""" -try: - result = await maybe_fail(True) -except Exception as e: - print(f"Second call failed: {type(e).__name__}") -""") - assert "Second call failed:" in str(output2) - - # Third call succeeds - tool still works - output3 = interpreter.execute(""" -result = await maybe_fail(False) -print(f"Third call: {result}") -""") - assert "Third call: success" in str(output3) - assert call_count == 3 - finally: - interpreter.shutdown() - - def test_parallel_calls_with_some_failures(self): - """asyncio.gather handles mix of successful and failing tool calls.""" - - def conditional_tool(item_id: int) -> str: - if item_id % 3 == 0: - raise ValueError(f"Item {item_id} failed") - return f"Item {item_id} ok" - - interpreter = JspiBackend( - preinstall_packages=False, tools={"conditional_tool": conditional_tool} - ) - try: - output = interpreter.execute(""" -import asyncio - -async def safe_call(item_id): - try: - return await conditional_tool(item_id) - except Exception as e: - return f"Error: {e}" - -# Run 10 parallel calls (items 0,3,6,9 will fail) -tasks = [safe_call(i) for i in range(10)] -results = await asyncio.gather(*tasks) - -successes = [r for r in results if r.startswith("Item") and "ok" in r] -failures = [r for r in results if r.startswith("Error")] - -print(f"Successes: {len(successes)}") -print(f"Failures: {len(failures)}") -print(f"Total: {len(results)}") -""") - assert "Successes: 6" in str(output) # 1,2,4,5,7,8 - assert "Failures: 4" in str(output) # 0,3,6,9 - assert "Total: 10" in str(output) - finally: - interpreter.shutdown() - - def test_tool_timeout_simulation(self): - """Tools that take too long can be handled with asyncio.wait_for.""" - import time - - def slow_tool(delay: float) -> str: - time.sleep(delay) - return f"Completed after {delay}s" - - interpreter = JspiBackend(preinstall_packages=False, tools={"slow_tool": slow_tool}) - try: - output = interpreter.execute(""" -import asyncio - -async def with_timeout(delay, timeout): - try: - result = await asyncio.wait_for(slow_tool(delay), timeout=timeout) - return f"OK: {result}" - except asyncio.TimeoutError: - return "TIMEOUT" - -# Fast call should succeed -result1 = await with_timeout(0.01, timeout=2.0) -print(f"Fast call: {result1}") - -# Note: actual timeout test would be slow, just verify the pattern works -print("Timeout pattern works") -""") - assert "Fast call: OK:" in str(output) - assert "Timeout pattern works" in str(output) - finally: - interpreter.shutdown() - - def test_tool_returns_none(self): - """Tools that return None round-trip as Python None.""" - - def none_tool() -> None: - return None - - interpreter = JspiBackend(preinstall_packages=False, tools={"none_tool": none_tool}) - try: - output = interpreter.execute(""" -result = await none_tool() -print(f"Result is None: {result is None}") -print(f"Result type: {type(result).__name__}") -""") - assert "Result is None: True" in str(output) - assert "Result type: NoneType" in str(output) - finally: - interpreter.shutdown() - - def test_tool_returns_empty_values(self): - """Tools returning empty strings/lists/dicts are handled correctly.""" - - def empty_string() -> str: - return "" - - def empty_list() -> list: - return [] - - def empty_dict() -> dict: - return {} - - interpreter = JspiBackend( - preinstall_packages=False, - tools={ - "empty_string": empty_string, - "empty_list": empty_list, - "empty_dict": empty_dict, - }, - ) - try: - output = interpreter.execute(""" -s = await empty_string() -l = await empty_list() -d = await empty_dict() - -print(f"String empty: {s == ''}") -print(f"List empty: {l == []}") -print(f"Dict empty: {d == dict()}") # Use dict() instead of {{}} for clarity -""") - assert "String empty: True" in str(output) - assert "List empty: True" in str(output) - assert "Dict empty: True" in str(output) - finally: - interpreter.shutdown() - - -class TestRealisticWorkloads: - """Tests simulating realistic API-like workloads with random delays.""" - - def test_random_latency_tool_calls(self): - """Tools with random latency (simulating network jitter).""" - import random - import time - - call_times = [] - - def api_call(item_id: int) -> dict: - # Random delay 1-10ms to simulate network latency (fast for tests) - delay = random.uniform(0.001, 0.01) - time.sleep(delay) - call_times.append((item_id, delay)) - return {"id": item_id, "status": "ok", "latency_ms": int(delay * 1000)} - - interpreter = JspiBackend(preinstall_packages=False, tools={"api_call": api_call}) - try: - output = interpreter.execute(""" -import asyncio - -# Make 10 parallel calls with random latency -tasks = [api_call(i) for i in range(10)] -results = await asyncio.gather(*tasks) - -# Verify all completed -completed = [r for r in results if r.get("status") == "ok"] -print(f"Completed: {len(completed)}/10") - -# Show latency range -latencies = [r["latency_ms"] for r in results] -print(f"Latency range: {min(latencies)}-{max(latencies)}ms") -""") - assert "Completed: 10/10" in str(output) - assert len(call_times) == 10 - finally: - interpreter.shutdown() - - def test_mixed_fast_and_slow_calls(self): - """Mix of fast and slow tool calls running concurrently.""" - import time - - def variable_speed(item_id: int, slow: bool) -> str: - delay = 0.01 if slow else 0.001 # 10ms vs 1ms - time.sleep(delay) - return f"item_{item_id}_{'slow' if slow else 'fast'}" - - interpreter = JspiBackend( - preinstall_packages=False, tools={"variable_speed": variable_speed} - ) - try: - output = interpreter.execute(""" -import asyncio - -# Mix of fast (even IDs) and slow (odd IDs) calls -tasks = [variable_speed(i, slow=(i % 2 == 1)) for i in range(8)] -results = await asyncio.gather(*tasks) - -fast_count = sum(1 for r in results if 'fast' in r) -slow_count = sum(1 for r in results if 'slow' in r) - -print(f"Fast: {fast_count}, Slow: {slow_count}") -print(f"All completed: {len(results) == 8}") -""") - assert "Fast: 4, Slow: 4" in str(output) - assert "All completed: True" in str(output) - finally: - interpreter.shutdown() - - def test_sequential_batches_with_delays(self): - """Multiple sequential batches of parallel calls.""" - import random - import time - - batch_results = [] - - def batch_item(batch_id: int, item_id: int) -> dict: - time.sleep(random.uniform(0.001, 0.005)) # 1-5ms - result = {"batch": batch_id, "item": item_id} - batch_results.append(result) - return result - - interpreter = JspiBackend( - preinstall_packages=False, tools={"batch_item": batch_item} - ) - try: - output = interpreter.execute(""" -import asyncio - -all_results = [] - -# Run 3 sequential batches, each with 5 parallel calls -for batch_id in range(3): - tasks = [batch_item(batch_id, i) for i in range(5)] - batch_results = await asyncio.gather(*tasks) - all_results.extend(batch_results) - print(f"Batch {batch_id} done: {len(batch_results)} items") - -print(f"Total items: {len(all_results)}") - -# Verify all batches represented -batches = set(r["batch"] for r in all_results) -print(f"Unique batches: {sorted(batches)}") -""") - assert "Total items: 15" in str(output) - assert "Unique batches: [0, 1, 2]" in str(output) - assert len(batch_results) == 15 - finally: - interpreter.shutdown() - - def test_intermittent_failures_with_retry(self): - """Simulating flaky API with retries.""" - import random - import time - - call_count = {"total": 0, "failures": 0} - - def flaky_api(item_id: int) -> dict: - call_count["total"] += 1 - time.sleep(random.uniform(0.001, 0.003)) - # 30% chance of failure - if random.random() < 0.3: - call_count["failures"] += 1 - raise RuntimeError(f"Transient error for item {item_id}") - return {"id": item_id, "success": True} - - interpreter = JspiBackend(preinstall_packages=False, tools={"flaky_api": flaky_api}) - try: - output = interpreter.execute(""" -import asyncio -import random - -async def call_with_retry(item_id, max_retries=3): - for attempt in range(max_retries): - try: - return await flaky_api(item_id) - except Exception as e: - if attempt == max_retries - 1: - return {"id": item_id, "success": False, "error": str(e)} - await asyncio.sleep(0.01) # Brief delay before retry - -# Run 10 items with retry logic -tasks = [call_with_retry(i) for i in range(10)] -results = await asyncio.gather(*tasks) - -successes = sum(1 for r in results if r.get("success")) -failures = sum(1 for r in results if not r.get("success")) - -print(f"Successes: {successes}") -print(f"Final failures: {failures}") -print(f"All processed: {len(results) == 10}") -""") - assert "All processed: True" in str(output) - # With retries, most should succeed despite 30% failure rate - assert call_count["total"] >= 10 # At least 10 calls, likely more due to retries - finally: - interpreter.shutdown() - - def test_high_concurrency_stress(self): - """Stress test with many concurrent calls and random delays.""" - import random - import time - - call_log = [] - - def stress_call(call_id: int) -> str: - start = time.perf_counter() - time.sleep(random.uniform(0.001, 0.005)) - elapsed = time.perf_counter() - start - call_log.append({"id": call_id, "elapsed": elapsed}) - return f"call_{call_id}_done" - - interpreter = JspiBackend( - preinstall_packages=False, tools={"stress_call": stress_call} - ) - try: - output = interpreter.execute(""" -import asyncio - -# 30 concurrent calls -tasks = [stress_call(i) for i in range(30)] -results = await asyncio.gather(*tasks) - -completed = sum(1 for r in results if 'done' in r) -print(f"Completed: {completed}/30") -""") - assert "Completed: 30/30" in str(output) - assert len(call_log) == 30 - finally: - interpreter.shutdown() - - def test_tools_persist_after_stress(self): - """Verify tools still work after high-concurrency stress test.""" - import random - import time - - def stress_tool(item_id: int) -> str: - time.sleep(random.uniform(0.001, 0.005)) - return f"result_{item_id}" - - interpreter = JspiBackend( - preinstall_packages=False, tools={"stress_tool": stress_tool} - ) - try: - # First: stress test - output1 = interpreter.execute(""" -import asyncio -tasks = [stress_tool(i) for i in range(25)] -results = await asyncio.gather(*tasks) -print(f"Stress test: {len(results)} calls") -""") - assert "Stress test: 25 calls" in str(output1) - - # Second: verify tool still works - output2 = interpreter.execute(""" -result = await stress_tool(999) -print(f"After stress: {result}") -""") - assert "After stress: result_999" in str(output2) - - # Third: another batch to confirm - output3 = interpreter.execute(""" -import asyncio -tasks = [stress_tool(i) for i in range(5)] -results = await asyncio.gather(*tasks) -print(f"Final batch: {len(results)} calls") -""") - assert "Final batch: 5 calls" in str(output3) - finally: - interpreter.shutdown() - - -class TestPydanticSerialization: - """Tests for Pydantic model serialization in tool calls.""" - - def test_pydantic_model_as_tool_argument(self): - """Pydantic models passed to tools are serialized correctly.""" - received_data = [] - - def process_model(data: dict) -> str: - received_data.append(data) - return f"Processed: {data.get('name', 'unknown')}" - - interpreter = JspiBackend(tools={"process_model": process_model}) - try: - output = interpreter.execute(""" -from pydantic import BaseModel - -class UserInput(BaseModel): - name: str - age: int - tags: list[str] = [] - -user = UserInput(name="Alice", age=30, tags=["admin", "active"]) -result = await process_model(user) -print(result) -""") - assert "Processed: Alice" in str(output) - assert len(received_data) == 1 - assert received_data[0]["name"] == "Alice" - assert received_data[0]["age"] == 30 - assert received_data[0]["tags"] == ["admin", "active"] - finally: - interpreter.shutdown() - - def test_pydantic_model_tool_result_is_mapping(self): - from pydantic import BaseModel - - class Source(BaseModel): - title: str - - class Retrieval(BaseModel): - model_name: str - sources: list[Source] - - def retrieve() -> Retrieval: - return Retrieval( - model_name="test-index", - sources=[Source(title="PredictRLM documentation")], - ) - - interpreter = JspiBackend(tools={"retrieve": retrieve}) - try: - output = interpreter.execute(""" -result = await retrieve() -print(result["model_name"]) -print(result["sources"][0]["title"]) -""") - assert "test-index" in str(output) - assert "PredictRLM documentation" in str(output) - finally: - interpreter.shutdown() - - def test_nested_pydantic_tool_results_are_mappings(self): - from pydantic import BaseModel - - class Source(BaseModel): - title: str - - def search() -> dict: - return {"results": [Source(title="PredictRLM documentation")]} - - interpreter = JspiBackend(tools={"search": search}) - try: - output = interpreter.execute(""" -result = await search() -print(result["results"][0]["title"]) -""") - assert "PredictRLM documentation" in str(output) - finally: - interpreter.shutdown() - - def test_nested_pydantic_models(self): - """Nested Pydantic models are serialized correctly.""" - received_data = [] - - def process_nested(data: dict) -> str: - received_data.append(data) - return f"Got {len(data.get('items', []))} items" - - interpreter = JspiBackend(tools={"process_nested": process_nested}) - try: - output = interpreter.execute(""" -from pydantic import BaseModel -from typing import List, Optional - -class Item(BaseModel): - id: int - name: str - price: float - -class Order(BaseModel): - order_id: str - customer: str - items: List[Item] - notes: Optional[str] = None - -order = Order( - order_id="ORD-123", - customer="Bob", - items=[ - Item(id=1, name="Widget", price=9.99), - Item(id=2, name="Gadget", price=19.99), - ], - notes="Rush delivery" -) - -result = await process_nested(order) -print(result) -""") - assert "Got 2 items" in str(output) - assert len(received_data) == 1 - assert received_data[0]["order_id"] == "ORD-123" - assert len(received_data[0]["items"]) == 2 - assert received_data[0]["items"][0]["name"] == "Widget" - finally: - interpreter.shutdown() - - def test_list_of_pydantic_models(self): - """Lists of Pydantic models are serialized correctly.""" - received_data = [] - - def process_list(items: list) -> str: - received_data.append(items) - return f"Received {len(items)} items" - - interpreter = JspiBackend(tools={"process_list": process_list}) - try: - output = interpreter.execute(""" -from pydantic import BaseModel - -class Task(BaseModel): - title: str - done: bool = False - -tasks = [ - Task(title="Task 1", done=True), - Task(title="Task 2"), - Task(title="Task 3", done=True), -] - -result = await process_list(tasks) -print(result) -""") - assert "Received 3 items" in str(output) - assert len(received_data) == 1 - assert len(received_data[0]) == 3 - assert received_data[0][0]["title"] == "Task 1" - assert received_data[0][0]["done"] is True - assert received_data[0][1]["done"] is False - finally: - interpreter.shutdown() - - def test_pydantic_in_kwargs(self): - """Pydantic models passed as kwargs are serialized correctly.""" - received_kwargs = [] - - def with_kwargs(**kwargs) -> str: - received_kwargs.append(kwargs) - return f"Got kwargs: {list(kwargs.keys())}" - - interpreter = JspiBackend(tools={"with_kwargs": with_kwargs}) - try: - output = interpreter.execute(""" -from pydantic import BaseModel - -class Config(BaseModel): - debug: bool = False - max_retries: int = 3 - -config = Config(debug=True, max_retries=5) -result = await with_kwargs(name="test", config=config, count=10) -print(result) -""") - assert "Got kwargs:" in str(output) - assert len(received_kwargs) == 1 - assert received_kwargs[0]["name"] == "test" - assert received_kwargs[0]["config"]["debug"] is True - assert received_kwargs[0]["config"]["max_retries"] == 5 - finally: - interpreter.shutdown() - - def test_tools_work_after_pydantic_serialization(self): - """Tools continue working after Pydantic model serialization.""" - - def echo(data: dict) -> dict: - return {"echoed": data} - - interpreter = JspiBackend(tools={"echo": echo}) - try: - # First call with Pydantic model - output1 = interpreter.execute(""" -from pydantic import BaseModel - -class MyModel(BaseModel): - value: int - -m = MyModel(value=42) -result = await echo(m) -print(f"First: {result}") -""") - assert "echoed" in str(output1) - - # Second call - tool should still work - output2 = interpreter.execute(""" -result = await echo({"plain": "dict"}) -print(f"Second: {result}") -""") - assert "plain" in str(output2) - - # Third call with another Pydantic model - output3 = interpreter.execute(""" -from pydantic import BaseModel - -class AnotherModel(BaseModel): - items: list[str] - -m2 = AnotherModel(items=["a", "b", "c"]) -result = await echo(m2) -print(f"Third: {result}") -""") - assert "items" in str(output3) - finally: - interpreter.shutdown() - - def test_non_serializable_falls_back_to_string(self): - """Non-serializable objects are converted to strings gracefully.""" - received_args = [] - - def receive_anything(x) -> str: - received_args.append(x) - return f"Got type: {type(x).__name__}" - - interpreter = JspiBackend(tools={"receive_anything": receive_anything}) - try: - # Pass a coroutine (non-standard object) - should be converted to string - output = interpreter.execute(""" -async def dummy(): - return 42 - -coro = dummy() -result = await receive_anything(coro) -print(result) -# Clean up -coro.close() -""") - # Should succeed with string representation - assert "Got type: str" in str(output) - assert len(received_args) == 1 - # The coroutine should have been converted to its string repr - assert "coroutine" in str(received_args[0]).lower() - - # Tool should still work normally - output2 = interpreter.execute(""" -result = await receive_anything({"normal": "dict"}) -print(result) -""") - assert "Got type: dict" in str(output2) - finally: - interpreter.shutdown() - - -class TestNoneValueSerialization: - """Tests that None/True/False in Pydantic models survive sandbox injection.""" - - def test_pydantic_model_with_none_fields_injected_as_variable(self): - """Pydantic models with None fields are accessible in the sandbox.""" - from pydantic import BaseModel, Field - - class ExtractedItem(BaseModel): + class ExtractedItem(BaseModel): title: str priority: str | None = Field(default=None) - due_date: str | None = Field(default=None) - active: bool = True - - items = [ - ExtractedItem(title="Task A", priority=None, due_date=None, active=True), - ExtractedItem(title="Task B", priority="high", due_date="2025-01-01", active=False), - ] - - interpreter = JspiBackend(preinstall_packages=False) - try: - result = interpreter.execute( - """ -print(f"count={len(items)}") -print(f"a_priority={items[0]['priority']}") -print(f"a_due={items[0]['due_date']}") -print(f"a_active={items[0]['active']}") -print(f"b_priority={items[1]['priority']}") -print(f"b_active={items[1]['active']}") -""", - variables={"items": items}, - ) - output = str(result) - assert "count=2" in output - assert "a_priority=None" in output - assert "a_due=None" in output - assert "a_active=True" in output - assert "b_priority=high" in output - assert "b_active=False" in output - finally: - interpreter.shutdown() - - def test_plain_dict_with_none_values_injected_as_variable(self): - """Plain dicts with None values are accessible in the sandbox.""" - data = [ - {"name": "Alice", "email": None, "verified": True}, - {"name": "Bob", "email": "bob@test.com", "verified": False}, - ] - - interpreter = JspiBackend(preinstall_packages=False) - try: - result = interpreter.execute( - """ -print(f"alice_email={data[0]['email']}") -print(f"alice_verified={data[0]['verified']}") -print(f"bob_verified={data[1]['verified']}") -""", - variables={"data": data}, - ) - output = str(result) - assert "alice_email=None" in output - assert "alice_verified=True" in output - assert "bob_verified=False" in output - finally: - interpreter.shutdown() - - -class TestPredictRLMWithPydanticModels: - """Tests for using Pydantic models with PredictRLM.""" - - def test_predict_rlm_with_pydantic_model_containing_methods(self): - """Test that PredictRLM works with Pydantic models that have methods.""" - # This test verifies that the fix works by ensuring schemas with methods can be serialized - # The actual integration test with PredictRLM would require a more complex setup - - # Simple test to verify that our defensive serialization works - interpreter = JspiBackend(preinstall_packages=False) - try: - result = interpreter.execute(""" -from pydantic import BaseModel, Field - -class ExtractedItem(BaseModel): - category: str = Field(description="Category") - title: str = Field(description="Title") - - def custom_method(self): - return f"{self.category}: {self.title}" - - @property - def formatted(self): - return self.custom_method() - -# Create an instance to verify the model works -item = ExtractedItem(category="Test", title="Item 1") -print(f"Item: {item.category} - {item.title}") -print(f"Formatted: {item.formatted}") - -# Verify the schema can be extracted and is JSON-serializable -import json -schema = ExtractedItem.model_json_schema() -json.dumps(schema) # This should not raise an error -print("Schema serialization successful") -""") - - assert "Schema serialization successful" in str(result) - assert "Test - Item 1" in str(result) - - finally: - interpreter.shutdown() - - -class TestCustomPydanticTypesInSignatures: - """Tests for custom Pydantic types in predict() return signatures.""" - - def test_custom_pydantic_type_in_return_signature(self): - """Custom Pydantic type in return signature extracts schemas correctly.""" - received_schemas = [] - - def mock_predict(signature: str, pydantic_schemas=None, **kwargs): - received_schemas.append(pydantic_schemas) - # Return mock data matching the expected type - return {"tasks": [{"category": "Test", "title": "Task 1"}]} - - interpreter = JspiBackend(tools={"predict": mock_predict}) - try: - result = interpreter.execute(""" -from pydantic import BaseModel - -class TaskItem(BaseModel): - category: str - title: str - -# This should extract TaskItem schema and pass it to predict -result = await predict("text: str -> tasks: list[TaskItem]", text="test input") -print(f"Got {len(result['tasks'])} tasks") -print(f"First task: {result['tasks'][0]}") -""") - assert "Got 1 tasks" in str(result) - # Verify schema was extracted and passed - assert len(received_schemas) == 1 - schemas = received_schemas[0] - assert schemas is not None - assert "TaskItem" in schemas - # Verify schema structure - assert "properties" in schemas["TaskItem"] - assert "category" in schemas["TaskItem"]["properties"] - assert "title" in schemas["TaskItem"]["properties"] - finally: - interpreter.shutdown() - - def test_nested_pydantic_types_in_signature(self): - """Nested Pydantic models extract schemas with $defs for nested types.""" - received_schemas = [] - - def mock_predict(signature: str, pydantic_schemas=None, **kwargs): - received_schemas.append(pydantic_schemas) - return { - "person": {"name": "Alice", "address": {"street": "123 Main", "city": "NYC"}} - } - - interpreter = JspiBackend(tools={"predict": mock_predict}) - try: - result = interpreter.execute(""" -from pydantic import BaseModel - -class Address(BaseModel): - street: str - city: str - -class Person(BaseModel): - name: str - address: Address # Nested model - -result = await predict("text: str -> person: Person", text="test") -print(f"Got person: {result['person'].name} at {result['person'].address.city}") -""") - assert "Got person: Alice at NYC" in str(result) - # Verify schema was extracted - assert len(received_schemas) == 1 - schemas = received_schemas[0] - assert schemas is not None - assert "Person" in schemas - # Verify nested type is in $defs - assert "$defs" in schemas["Person"] - assert "Address" in schemas["Person"]["$defs"] - finally: - interpreter.shutdown() - - def test_optional_pydantic_type_in_signature(self): - """Optional custom Pydantic type in signature works correctly.""" - received_schemas = [] - - def mock_predict(signature: str, pydantic_schemas=None, **kwargs): - received_schemas.append(pydantic_schemas) - return {"item": {"name": "Widget", "price": 9.99}} - - interpreter = JspiBackend(tools={"predict": mock_predict}) - try: - result = interpreter.execute(""" -from pydantic import BaseModel -from typing import Optional - -class Item(BaseModel): - name: str - price: float - -result = await predict("text: str -> item: Optional[Item]", text="test") -print(f"Got item: {result['item'].name}") -""") - assert "Got item: Widget" in str(result) - # Verify schema was extracted - assert len(received_schemas) == 1 - schemas = received_schemas[0] - assert schemas is not None - assert "Item" in schemas - finally: - interpreter.shutdown() - - def test_builtin_types_not_included_in_schemas(self): - """Built-in types like str, int, bool are not included in pydantic_schemas.""" - received_schemas = [] - - def mock_predict(signature: str, pydantic_schemas=None, **kwargs): - received_schemas.append(pydantic_schemas) - return {"name": "test", "count": 42, "active": True} - - interpreter = JspiBackend(tools={"predict": mock_predict}) - try: - result = interpreter.execute(""" -result = await predict("text: str -> name: str, count: int, active: bool", text="test") -print(f"Got: {result['name']}, {result['count']}, {result['active']}") -""") - assert "Got: test, 42, True" in str(result) - # No pydantic_schemas should be passed for built-in types - assert len(received_schemas) == 1 - assert received_schemas[0] is None or len(received_schemas[0]) == 0 - finally: - interpreter.shutdown() - - def test_dspy_image_type_not_in_schemas(self): - """dspy.Image type is not included in pydantic_schemas (it's built-in).""" - received_schemas = [] - - def mock_predict(signature: str, pydantic_schemas=None, **kwargs): - received_schemas.append(pydantic_schemas) - return {"text": "extracted text"} - - interpreter = JspiBackend(tools={"predict": mock_predict}) - try: - result = interpreter.execute(""" -result = await predict("img: dspy.Image -> text: str", img="http://example.com/img.png") -print(f"Got: {result['text']}") -""") - assert "Got: extracted text" in str(result) - # dspy.Image should not be in schemas - assert len(received_schemas) == 1 - schemas = received_schemas[0] - assert schemas is None or "Image" not in schemas - finally: - interpreter.shutdown() - - def test_multiple_custom_types_in_signature(self): - """Multiple custom Pydantic types in same signature are all extracted.""" - received_schemas = [] - - def mock_predict(signature: str, pydantic_schemas=None, **kwargs): - received_schemas.append(pydantic_schemas) - return { - "header": {"title": "Doc", "version": "1.0"}, - "items": [{"name": "Item1", "qty": 1}], - } - - interpreter = JspiBackend(tools={"predict": mock_predict}) - try: - result = interpreter.execute(""" -from pydantic import BaseModel - -class Header(BaseModel): - title: str - version: str - -class LineItem(BaseModel): - name: str - qty: int - -result = await predict("text: str -> header: Header, items: list[LineItem]", text="test") -print(f"Got header: {result['header'].title}") -print(f"Got {len(result['items'])} items") -""") - assert "Got header: Doc" in str(result) - assert "Got 1 items" in str(result) - # Both types should be in schemas - assert len(received_schemas) == 1 - schemas = received_schemas[0] - assert schemas is not None - assert "Header" in schemas - assert "LineItem" in schemas - finally: - interpreter.shutdown() - - def test_schema_extraction_survives_tool_errors(self): - """Schema extraction works correctly even after tool errors.""" - call_count = {"count": 0} - - def mock_predict(signature: str, pydantic_schemas=None, **kwargs): - call_count["count"] += 1 - if call_count["count"] == 1: - raise RuntimeError("Intentional error") - return {"item": {"value": 42}} - - interpreter = JspiBackend(tools={"predict": mock_predict}) - try: - # First call fails - output1 = interpreter.execute(""" -from pydantic import BaseModel - -class MyType(BaseModel): - value: int - -try: - result = await predict("text: str -> item: MyType", text="test") -except Exception as e: - print(f"First call failed: {type(e).__name__}") -""") - assert "First call failed" in str(output1) - - # Second call succeeds with same type - output2 = interpreter.execute(""" -from pydantic import BaseModel - -class MyType(BaseModel): - value: int - -result = await predict("text: str -> item: MyType", text="test") -print(f"Second call got: {result['item'].value}") -""") - assert "Second call got: 42" in str(output2) - finally: - interpreter.shutdown() - - -class TestPydanticLiteralSerialization: - """Tests that Pydantic models with Literal types serialize correctly through the interpreter.""" - - def test_pydantic_model_with_literal_in_predict_signature(self): - """Pydantic models with Literal types should work in predict signatures.""" - - # Mock predict tool that echoes back the schema it received - async def mock_predict( - signature: str, *, instructions: str = None, pydantic_schemas: dict = None, **kwargs - ): - # Check if we received the schema - if pydantic_schemas and "TaskItem" in pydantic_schemas: - schema = pydantic_schemas["TaskItem"] - # Verify the schema has the expected structure - props = schema.get("properties", {}) - if "priority" in props: - # Extract priority anyOf to check for Literal values - priority_spec = props["priority"].get("anyOf", []) - # Return a mock task with the schema info - return { - "items": [ - { - "category": "Test Category", - "title": "Test Task", - "description": "Test Description", - "priority": "high", - "due_date": None, - "_schema_received": True, - "_priority_spec_count": len(priority_spec), - } - ] - } - return {"items": []} - - interpreter = JspiBackend(preinstall_packages=True, tools={"predict": mock_predict}) - - try: - # Execute code that defines a Pydantic model with Literal and uses it - result = interpreter.execute(""" -import asyncio -from typing import Literal, Optional -from pydantic import BaseModel, Field - -Priority = Optional[Literal["urgent", "high", "medium", "low"]] - -class TaskItem(BaseModel): - category: str = Field(description="Task category") - title: str - description: str - priority: Priority = None - due_date: Optional[str] = None - -# Call predict with the Pydantic model in the signature -result = await predict( - "page: str -> items: list[TaskItem]", - instructions="Extract tasks", - page="test page content" -) - -# Check result -print(f"Got {len(result['items'])} items") -if result['items']: - item = result['items'][0] - print(f"First item category: {item.category}") - print(f"Schema received: {getattr(item, '_schema_received', False)}") - print(f"Priority spec count: {getattr(item, '_priority_spec_count', 0)}") -""") - assert "Got 1 items" in str(result) - assert "First item category: Test Category" in str(result) - assert "Schema received: True" in str(result) - # Should have anyOf with 5 options (4 Literal values + null) - assert "Priority spec count: 2" in str(result) # anyOf with enum and null - - finally: - interpreter.shutdown() - - def test_pydantic_model_serialization_through_tool(self): - """Pydantic models should serialize correctly when passed to tools.""" - - received_data = [] - - async def process_items(items): - """Tool that receives a list of Pydantic models.""" - # Store what we received for verification - received_data.append(items) - return f"Processed {len(items)} items" - - interpreter = JspiBackend( - preinstall_packages=True, tools={"process_items": process_items} - ) - - try: - result = interpreter.execute(""" -from typing import Literal, Optional -from pydantic import BaseModel - -class Priority(BaseModel): - level: Literal["urgent", "high", "medium", "low"] - -class Task(BaseModel): - title: str - priority: Optional[Priority] = None - -# Create tasks with Literal values -tasks = [ - Task(title="Task 1", priority=Priority(level="high")), - Task(title="Task 2", priority=Priority(level="urgent")), - Task(title="Task 3"), # No priority -] - -# Pass to tool - models should be serialized to dicts -result = await process_items(tasks) -print(result) - -# Verify serialization worked -for i, task in enumerate(tasks): - if task.priority: - print(f"Task {i+1} priority: {task.priority.level}") - else: - print(f"Task {i+1} has no priority") -""") - assert "Processed 3 items" in str(result) - assert "Task 1 priority: high" in str(result) - assert "Task 2 priority: urgent" in str(result) - assert "Task 3 has no priority" in str(result) - - # Verify the tool received proper dicts - assert len(received_data) == 1 - items = received_data[0] - assert len(items) == 3 - assert items[0]["title"] == "Task 1" - assert items[0]["priority"]["level"] == "high" - assert items[2]["priority"] is None - - finally: - interpreter.shutdown() - - def test_dspy_history_serialization_with_rlm(self): - """Test that DSPy history serialization doesn't produce warnings after RLM run.""" - import warnings - from typing import Literal, Optional - - import dspy - from pydantic import BaseModel, Field - - from predict_rlm import PredictRLM - - # Define Pydantic models with Literal like in the notebook - class ExtractedItem(BaseModel): - """A task item extracted from documents.""" - - category: str = Field(description="Category from extraction") - title: str = Field(description="Task title") - description: str = Field(description="Task description") - priority: Optional[Literal["urgent", "high", "medium", "low"]] = Field( - default=None, description="Priority level" - ) - due_date: Optional[str] = Field(default=None, description="Due date") - doc_id: str = Field(description="Source document ID") - page: int = Field(description="Page number (0-indexed)") - - class ExtractionResult(BaseModel): - """All items extracted from documents.""" - - items: list[ExtractedItem] = Field(description="List of extracted items") - - # Signature using the Pydantic models - class ExtractFromDocuments(dspy.Signature): - """Extract structured items from documents.""" - - documents: list[dict] = dspy.InputField(desc="Documents to analyze") - prompt: str = dspy.InputField(desc="Extraction instructions") - result: ExtractionResult = dspy.OutputField( - desc="Task items with category, title, description" - ) - - # Mock tools - async def get_pages(doc_id: str, pages: list[int], format: str): - return [f"fake_image_{i}" for i in pages] - - async def search(query: str, doc_id: Optional[str] = None, limit: int = 5): - return [{"doc_id": "test", "page": 0, "text": "test", "relevance": 1.0}] - - # Use a real LM model (OpenAI) that would trigger the issue - test_lm = dspy.LM(model="openai/gpt-5-mini-2025-08-07", cache=False) - - # Create the RLM - rlm = PredictRLM( - ExtractFromDocuments, - max_iterations=2, - verbose=True, - tools={"get_pages": get_pages, "search": search}, - ) - - # Prepare test data like the notebook - test_docs = [{"doc_id": "doc1", "file_name": "test.pdf", "page_count": 2}] - test_prompt = "Extract actionable tasks from RFP" - - # Capture warnings during execution and history access - with warnings.catch_warnings(record=True) as w: - warnings.simplefilter("always") - - try: - # Run the RLM with the test LM context - with dspy.context(lm=test_lm): - # This might fail due to incomplete setup, which is fine - try: - _ = rlm(documents=test_docs, prompt=test_prompt) - except Exception as e: - # We expect some failures, but not serialization warnings - print(f"Expected failure: {e}") - - # The critical part: access history which triggers serialization - # This is what happens in the notebook cell - _ = list(test_lm.history) - - except Exception as e: - # If we get serialization warnings, they should be captured - print(f"Exception during test: {e}") - - # Check for Pydantic serialization warnings from OUR code only. - # DSPy internals (cache.py, base_lm.py) emit these warnings when - # serializing history entries with Literal types — that's upstream. - serialization_warnings = [ - warning - for warning in w - if "PydanticSerializationUnexpectedValue" in str(warning.message) - and "predict_rlm" in str(warning.filename) - ] - - if len(serialization_warnings) > 0: - print( - f"Got {len(serialization_warnings)} Pydantic serialization warnings from our code:" - ) - for warning in serialization_warnings: - print(f" - {warning.message}") - print(f" File: {warning.filename}:{warning.lineno}") - - assert len(serialization_warnings) == 0, ( - f"Got {len(serialization_warnings)} Pydantic serialization warnings from our code" - ) - - -class TestSerializationFailureRecovery: - """Tests that tools survive serialization failures (the original bug). - - The original issue: when sandbox code tries to pass a non-existent object - to a tool, the serialization fails and corrupts the global state, making - tools unavailable for subsequent calls. - - Our fix stores tools in a persistent `_repl_tools` module and re-injects - them before each execution, preventing state corruption from affecting - tool availability. - """ - - def test_pydantic_model_with_method_serialization(self): - """Test that Pydantic models with methods can be serialized properly when passed to predict().""" - - # Track what gets passed to predict - predict_calls = [] - - def mock_predict(signature: str, pydantic_schemas=None, **kwargs): - predict_calls.append( - {"signature": signature, "schemas": pydantic_schemas, "kwargs": kwargs} - ) - # Return mock items based on signature - return { - "items": [ - {"category": "Test", "title": "Item 1", "description": "Test description"} - ] - } - - interpreter = JspiBackend(tools={"predict": mock_predict}) - try: - # This should work without serialization errors - result = interpreter.execute(""" -from pydantic import BaseModel, Field - -class ExtractedItem(BaseModel): - category: str = Field(description="Category from the extraction prompt") - title: str = Field(description="Task title - concise, actionable") - description: str = Field(description="Task directive using verbatim wording") - - # This method should not cause serialization issues - def to_dict(self): - return self.model_dump() - - # This shouldn't either - @property - def formatted_title(self): - return f"[{self.category}] {self.title}" - -# Using the model in predict signature - this was failing with method serialization -result = await predict( - "page: dspy.Image -> items: list[ExtractedItem]", - instructions="Extract tasks from the page", - page="dummy_image_url" -) -print(f"Got {len(result['items'])} items") -print(f"First item category: {result['items'][0].category}") -""") - - # Verify it worked - assert "Got 1 items" in str(result) - assert "First item category: Test" in str(result) - - # Verify predict was called with proper schemas - assert len(predict_calls) == 1 - call = predict_calls[0] - assert "ExtractedItem" in str(call["schemas"]) if call["schemas"] else False - - # Tool should still work after this - result2 = interpreter.execute(""" -result = await predict("text: str -> answer: str", text="hello") -print("Predict still works!") -""") - assert "Predict still works!" in str(result2) - - finally: - interpreter.shutdown() - - def test_pydantic_model_json_schema_with_method_serialization(self): - """Test that model_json_schema() with methods doesn't cause JSON serialization errors.""" - - interpreter = JspiBackend(preinstall_packages=False) - try: - # This reproduces the exact issue from the user's example - result = interpreter.execute(""" -from pydantic import BaseModel, Field -import json - -class ExtractedItem(BaseModel): - category: str = Field(description="Category") - title: str = Field(description="Title") - - def custom_method(self): - '''Custom method that should not cause issues''' - return f"{self.category}: {self.title}" - - @property - def formatted(self): - '''Property that should not cause issues''' - return self.custom_method() - -# Get the JSON schema - this should work -schema = ExtractedItem.model_json_schema() - -# This was causing "TypeError: Object of type method is not JSON serializable" -# when the schema contains references to methods -try: - serialized = json.dumps(schema) - print("SUCCESS: Schema serialized without error") - print(f"Schema keys: {list(schema.keys())}") -except TypeError as e: - print(f"ERROR: {e}") - print("This is the bug we're testing for") -""") - - # The test passes if we can serialize the schema without errors - assert "SUCCESS: Schema serialized" in str(result) - assert "ERROR" not in str(result) - - finally: - interpreter.shutdown() - - def test_pydantic_schema_extraction_with_methods_in_predict(self): - """Test that the _get_pydantic_schemas function works with models containing methods.""" - # This test directly tests the schema extraction that happens in runner.js - received_schemas = [] - - def mock_predict(signature: str, pydantic_schemas=None, **kwargs): - import json - - # Try to serialize the schemas - this is where the error would occur - if pydantic_schemas: - try: - json.dumps(pydantic_schemas) - received_schemas.append({"success": True, "schemas": pydantic_schemas}) - except TypeError as e: - received_schemas.append({"success": False, "error": str(e)}) - return {"items": []} - - interpreter = JspiBackend(tools={"predict": mock_predict}) - try: - result = interpreter.execute(""" -from pydantic import BaseModel, Field - -class ExtractedItem(BaseModel): - category: str = Field(description="Category") - title: str = Field(description="Title") - - def to_dict(self): - return self.model_dump() - - @property - def formatted_title(self): - return f"[{self.category}] {self.title}" - -# This should extract the schema and pass it through -result = await predict("text: str -> items: list[ExtractedItem]", text="test") -print("Predict call completed") -""") - - assert "Predict call completed" in str(result) - assert len(received_schemas) == 1 - - # Check if schema was successfully serialized - schema_result = received_schemas[0] - if not schema_result["success"]: - print(f"Schema serialization failed: {schema_result['error']}") - assert schema_result["success"], ( - f"Schema serialization failed: {schema_result.get('error')}" - ) - - finally: - interpreter.shutdown() - - def test_reproduce_user_exact_error_scenario(self): - """Reproduce the exact error scenario from the user's example with ExtractedItem.""" - interpreter = JspiBackend(preinstall_packages=False) - try: - # This is the exact pattern from the user's error output - result = interpreter.execute(""" -from pydantic import BaseModel, Field -import json - -class ExtractedItem(BaseModel): - category: str = Field(description="Category from the extraction prompt") - title: str = Field(description="Task title - concise, actionable") - description: str = Field(description="Task directive using verbatim wording") - doc_id: str = Field(description="Source document ID") - page: int = Field(description="Page number where found (0-indexed)") - due_date: str | None = Field(default=None, description="Due date if explicitly stated") - priority: str | None = Field(default=None, description="Priority level") - - # Methods that should not affect schema serialization - def to_dict(self): - return self.model_dump() - - @property - def formatted(self): - return f"[{self.category}] {self.title}" - -# Get schema - this should work now with our fix -schema = ExtractedItem.model_json_schema() - -# Try to serialize it (this was failing before) -try: - serialized = json.dumps(schema) - print("SUCCESS: Schema serialized without error") - print(f"Schema has {len(schema.get('properties', {}))} properties") -except TypeError as e: - print(f"ERROR: {e}") - print("The JSON serialization failed") -""") - - # Should succeed with our fix - assert "SUCCESS: Schema serialized" in str(result) - assert "ERROR" not in str(result) - - finally: - interpreter.shutdown() - - def test_pydantic_model_schema_with_function_validators(self): - """Test that Pydantic models with validators don't cause serialization issues.""" - received_schemas = [] - - def mock_predict(signature: str, pydantic_schemas=None, **kwargs): - import json - - if pydantic_schemas: - try: - # Check if we can JSON serialize the schema - json.dumps(pydantic_schemas) - received_schemas.append(pydantic_schemas) - except TypeError as e: - raise RuntimeError(f"Schema serialization failed: {e}") - return {"items": []} - - interpreter = JspiBackend(tools={"predict": mock_predict}) - try: - # Use field validators which might add function references to the schema - result = interpreter.execute(""" -from pydantic import BaseModel, Field, field_validator - -class ExtractedItem(BaseModel): - category: str = Field(description="Category") - title: str = Field(description="Title") - - @field_validator('category') - @classmethod - def validate_category(cls, v): - if not v: - raise ValueError('Category cannot be empty') - return v.upper() - - @field_validator('title') - @classmethod - def validate_title(cls, v): - return v.strip() - -# Extract the schema and use it in predict -result = await predict("text: str -> items: list[ExtractedItem]", text="test") -print("Schema extraction and serialization successful") -""") - - assert "Schema extraction and serialization successful" in str(result) - assert len(received_schemas) == 1 - - finally: - interpreter.shutdown() - - def test_undefined_variable_in_tool_arg(self): - """Tools survive when code passes an undefined variable.""" - call_log = [] - - def my_tool(data) -> str: - call_log.append(data) - return f"Got: {data}" - - interpreter = JspiBackend(preinstall_packages=False, tools={"my_tool": my_tool}) - try: - # First call - success - output1 = interpreter.execute(""" -result = await my_tool("valid") -print(result) -""") - assert "Got: valid" in str(output1) - assert len(call_log) == 1 - - # Second call - reference undefined variable (NameError) - # This should fail but NOT corrupt the tools - output2 = interpreter.execute(""" -try: - result = await my_tool(undefined_variable) - print("Should not reach here") -except NameError as e: - print(f"Caught NameError: {e}") -""") - assert "Caught NameError" in str(output2) - - # Third call - tool should still work! - output3 = interpreter.execute(""" -result = await my_tool("after_error") -print(result) -""") - assert "Got: after_error" in str(output3) - assert len(call_log) == 2 # Only 2 successful calls - finally: - interpreter.shutdown() - - def test_undefined_class_instantiation_in_tool_arg(self): - """Tools survive when code tries to instantiate undefined class.""" - call_log = [] - - def predict(signature: str, **kwargs) -> dict: - call_log.append({"sig": signature, "kwargs": kwargs}) - return {"output": f"Processed {len(kwargs)} args"} - - interpreter = JspiBackend(preinstall_packages=False, tools={"predict": predict}) - try: - # First call - success - output1 = interpreter.execute(""" -result = await predict("question -> answer", question="test") -print(result) -""") - assert "Processed 1 args" in str(output1) - - # Try to pass instance of undefined class - output2 = interpreter.execute(""" -try: - # NonExistentModel doesn't exist - this will raise NameError - obj = NonExistentModel(field="value") - result = await predict("data -> output", data=obj) - print("Should not reach here") -except NameError as e: - print(f"Caught: NameError") -""") - assert "Caught: NameError" in str(output2) - - # Tool should still work - output3 = interpreter.execute(""" -result = await predict("x -> y", x="recovery test") -print(result) -""") - assert "Processed 1 args" in str(output3) - assert len(call_log) == 2 - finally: - interpreter.shutdown() - - def test_serialization_error_mid_gather(self): - """Tools survive when one of many parallel calls has serialization error.""" - call_log = [] - - def process(item_id: int, data) -> str: - call_log.append({"id": item_id, "data": data}) - return f"Processed {item_id}" - - interpreter = JspiBackend(preinstall_packages=False, tools={"process": process}) - try: - # Run parallel calls where one will fail due to undefined reference - output = interpreter.execute(""" -import asyncio - -async def safe_process(item_id, data_factory): - try: - data = data_factory() - return await process(item_id, data) - except Exception as e: - return f"Error {item_id}: {type(e).__name__}" - -# Factory functions - one will fail -def valid1(): return "data1" -def valid2(): return "data2" -def invalid(): return undefined_var # Will raise NameError -def valid3(): return "data3" - -tasks = [ - safe_process(1, valid1), - safe_process(2, valid2), - safe_process(3, invalid), # This one fails - safe_process(4, valid3), -] -results = await asyncio.gather(*tasks) - -successes = [r for r in results if r.startswith("Processed")] -errors = [r for r in results if r.startswith("Error")] - -print(f"Successes: {len(successes)}") -print(f"Errors: {len(errors)}") -""") - assert "Successes: 3" in str(output) - assert "Errors: 1" in str(output) - - # Tool should still work after the mixed batch - output2 = interpreter.execute(""" -result = await process(99, "final") -print(result) -""") - assert "Processed 99" in str(output2) - finally: - interpreter.shutdown() - - def test_attribute_error_on_nonexistent_method(self): - """Tools survive when code calls non-existent method on result.""" - - def get_data() -> dict: - return {"items": [1, 2, 3]} - - interpreter = JspiBackend(preinstall_packages=False, tools={"get_data": get_data}) - try: - # First call - success - output1 = interpreter.execute(""" -data = await get_data() -print(f"Got {len(data['items'])} items") -""") - assert "Got 3 items" in str(output1) - - # Call method that doesn't exist on the result - output2 = interpreter.execute(""" -try: - data = await get_data() - # Try to call non-existent method - result = data.nonexistent_method() - print("Should not reach") -except AttributeError as e: - print("Caught AttributeError") -""") - assert "Caught AttributeError" in str(output2) - - # Tool should still work - output3 = interpreter.execute(""" -data = await get_data() -print(f"Recovery: {data['items'][0]}") -""") - assert "Recovery: 1" in str(output3) - finally: - interpreter.shutdown() - - def test_type_error_during_serialization(self): - """Tools survive TypeError during argument serialization.""" - received = [] - - def accept_dict(d: dict) -> str: - received.append(d) - return f"Got dict with {len(d)} keys" - - interpreter = JspiBackend( - preinstall_packages=False, tools={"accept_dict": accept_dict} - ) - try: - # Success - output1 = interpreter.execute(""" -result = await accept_dict({"key": "value"}) -print(result) -""") - assert "Got dict with 1 keys" in str(output1) - - # Try to pass something that looks like it has model_dump but fails - output2 = interpreter.execute(""" -class FakeModel: - def model_dump(self): - raise TypeError("Cannot serialize") - -try: - obj = FakeModel() - result = await accept_dict(obj) - print("Should not reach") -except Exception as e: - print(f"Caught: {type(e).__name__}") -""") - # Should catch the error (either TypeError or the wrapped error) - assert "Caught:" in str(output2) - - # Tool should still work - output3 = interpreter.execute(""" -result = await accept_dict({"recovery": True}) -print(result) -""") - assert "Got dict with 1 keys" in str(output3) - finally: - interpreter.shutdown() - - def test_tools_persist_through_repeated_failures(self): - """Tools remain available through multiple consecutive failures.""" - call_count = {"success": 0} - - def counter_tool(value: str) -> str: - call_count["success"] += 1 - return f"Count: {call_count['success']}" - - interpreter = JspiBackend( - preinstall_packages=False, tools={"counter_tool": counter_tool} - ) - try: - # Initial success - interpreter.execute(""" -result = await counter_tool("first") -print(result) -""") - assert call_count["success"] == 1 - - # Multiple failures in a row - for i in range(5): - interpreter.execute(f""" -try: - result = await counter_tool(undefined_var_{i}) -except NameError: - pass -""") - - # Tool should STILL work after 5 consecutive failures - output = interpreter.execute(""" -result = await counter_tool("recovery") -print(result) -""") - assert "Count: 2" in str(output) - assert call_count["success"] == 2 - finally: - interpreter.shutdown() - - def test_pydantic_model_with_missing_field_reference(self): - """Tools survive when Pydantic model references undefined field.""" - received = [] - - def process_model(data) -> str: - received.append(data) - return "processed" - - interpreter = JspiBackend(tools={"process_model": process_model}) - try: - # Success with valid model - output1 = interpreter.execute(""" -from pydantic import BaseModel - -class ValidModel(BaseModel): - name: str - value: int - -m = ValidModel(name="test", value=42) -result = await process_model(m) -print(f"First: {result}") -""") - assert "First: processed" in str(output1) - - # Try to create model referencing undefined variable - output2 = interpreter.execute(""" -from pydantic import BaseModel - -class AnotherModel(BaseModel): - data: str - -try: - # undefined_value doesn't exist - m = AnotherModel(data=undefined_value) - result = await process_model(m) - print("Should not reach") -except NameError: - print("Caught NameError as expected") -""") - assert "Caught NameError" in str(output2) - - # Tool should still work - output3 = interpreter.execute(""" -from pydantic import BaseModel - -class YetAnotherModel(BaseModel): - status: str - -m = YetAnotherModel(status="recovered") -result = await process_model(m) -print(f"Third: {result}") -""") - assert "Third: processed" in str(output3) - assert len(received) == 2 # Only successful calls - finally: - interpreter.shutdown() - - def test_tool_in_module_survives_serialization_crash(self): - """Verify tool remains in _repl_tools module after serialization error.""" - - def my_tool(x) -> str: - return f"got {x}" - - interpreter = JspiBackend(preinstall_packages=False, tools={"my_tool": my_tool}) - try: - # Cause a serialization failure inside the tool - output1 = interpreter.execute(""" -import sys - -class CrashingModel: - def model_dump(self): - raise RuntimeError("Serialization crashed!") - -try: - obj = CrashingModel() - await my_tool(obj) -except Exception as e: - print(f"Tool call failed: {type(e).__name__}") - -# Verify tool is still in _repl_tools module -_repl_tools = sys.modules.get('_repl_tools') -tool_in_module = hasattr(_repl_tools, 'my_tool') if _repl_tools else False -print(f"Tool in _repl_tools: {tool_in_module}") - -# Verify tool is still in globals -tool_in_globals = 'my_tool' in dir() -print(f"Tool in globals: {tool_in_globals}") -""") - assert "Tool call failed:" in str(output1) - assert "Tool in _repl_tools: True" in str(output1) - assert "Tool in globals: True" in str(output1) - - # Tool should work normally after the crash - output2 = interpreter.execute(""" -result = await my_tool("works!") -print(result) -""") - assert "got works!" in str(output2) - finally: - interpreter.shutdown() - - def test_tool_recovery_after_global_deletion(self): - """Tools can be recovered from _repl_tools even if deleted from globals. - - This simulates the worst case where globals get corrupted/cleared. - The re-injection mechanism should restore the tools. - """ - call_log = [] - - def track_tool(msg: str) -> str: - call_log.append(msg) - return f"tracked: {msg}" - - interpreter = JspiBackend( - preinstall_packages=False, tools={"track_tool": track_tool} - ) - try: - # First call works - output1 = interpreter.execute(""" -result = await track_tool("first") -print(result) -""") - assert "tracked: first" in str(output1) - - # Deliberately delete the tool from globals (simulating corruption) - output2 = interpreter.execute(""" -# Delete tool from globals -del track_tool - -# Verify it's gone from globals -tool_in_globals = 'track_tool' in dir() -print(f"After del - tool in globals: {tool_in_globals}") - -# But it should still be in _repl_tools module -import sys -_repl_tools = sys.modules.get('_repl_tools') -tool_in_module = hasattr(_repl_tools, 'track_tool') if _repl_tools else False -print(f"After del - tool in module: {tool_in_module}") -""") - assert "After del - tool in globals: False" in str(output2) - assert "After del - tool in module: True" in str(output2) - - # Next execution should re-inject the tool (our fix!) - output3 = interpreter.execute(""" -# Tool should be available again due to re-injection -result = await track_tool("after deletion") -print(result) -""") - assert "tracked: after deletion" in str(output3) - assert len(call_log) == 2 - finally: - interpreter.shutdown() - - def test_multiple_tools_survive_partial_corruption(self): - """All tools survive even when one tool's call crashes.""" - tool_a_calls = [] - tool_b_calls = [] - tool_c_calls = [] - - def tool_a(x) -> str: - tool_a_calls.append(x) - return f"A: {x}" - - def tool_b(x) -> str: - tool_b_calls.append(x) - return f"B: {x}" - - def tool_c(x) -> str: - tool_c_calls.append(x) - return f"C: {x}" - - interpreter = JspiBackend( - preinstall_packages=False, - tools={ - "tool_a": tool_a, - "tool_b": tool_b, - "tool_c": tool_c, - }, - ) - try: - # Use all tools successfully - output1 = interpreter.execute(""" -a = await tool_a("first") -b = await tool_b("first") -c = await tool_c("first") -print(f"{a}, {b}, {c}") -""") - assert "A: first" in str(output1) - assert "B: first" in str(output1) - assert "C: first" in str(output1) - - # Crash tool_b with a serialization error - output2 = interpreter.execute(""" -class BadArg: - def model_dump(self): - raise ValueError("Boom!") - -try: - await tool_b(BadArg()) -except Exception: - print("tool_b crashed") -""") - assert "tool_b crashed" in str(output2) - - # ALL tools should still work - output3 = interpreter.execute(""" -a = await tool_a("after") -b = await tool_b("after") -c = await tool_c("after") -print(f"{a}, {b}, {c}") -""") - assert "A: after" in str(output3) - assert "B: after" in str(output3) - assert "C: after" in str(output3) - - # Verify call counts - assert len(tool_a_calls) == 2 - assert len(tool_b_calls) == 2 # 2 successful, 1 failed (not counted) - assert len(tool_c_calls) == 2 - finally: - interpreter.shutdown() - - -class TestToolSurvivalAfterAsyncErrors: - """Test that tools survive when async code crashes with pending tool calls. - - This addresses the issue where tools would disappear after an error during - concurrent async execution (e.g., asyncio.gather with many predict calls). - The bug was that the response reader could consume the next command from - the interpreter if not properly stopped on error. - """ - - def test_tools_survive_after_unhandled_async_error(self): - """Tools remain available after an unhandled exception in async code.""" - - call_counts = {"slow": 0, "fast": 0} - - def slow_tool(x: str) -> str: - call_counts["slow"] += 1 - return f"slow: {x}" - - def fast_tool(x: str) -> str: - call_counts["fast"] += 1 - return f"fast: {x}" - - interpreter = JspiBackend( - preinstall_packages=False, - tools={"slow_tool": slow_tool, "fast_tool": fast_tool}, - ) - try: - # First, verify tools work normally - output1 = interpreter.execute(""" -r1 = await slow_tool("first") -r2 = await fast_tool("first") -print(f"Results: {r1}, {r2}") -""") - assert "slow: first" in str(output1) - assert "fast: first" in str(output1) - assert call_counts["slow"] == 1 - assert call_counts["fast"] == 1 - - # Now cause an error - but catch it so we can continue - output2 = interpreter.execute(""" -import asyncio - -async def fail_after_tool(): - r = await slow_tool("before_fail") - print(f"Got: {r}") - # Raise after the tool call - raise ValueError("Boom!") - -try: - await fail_after_tool() -except ValueError as e: - print(f"Caught: {e}") -""") - assert "Got: slow: before_fail" in str(output2) - assert "Caught: Boom!" in str(output2) - assert call_counts["slow"] == 2 - - # CRITICAL: Tools should still be available after the error - output3 = interpreter.execute(""" -# Tools should be defined and callable -r1 = await slow_tool("after_error") -r2 = await fast_tool("after_error") -print(f"After error: {r1}, {r2}") -""") - assert "slow: after_error" in str(output3) - assert "fast: after_error" in str(output3) - assert call_counts["slow"] == 3 - assert call_counts["fast"] == 2 - - finally: - interpreter.shutdown() - - def test_tools_survive_after_import_re(self): - """Importing 're' module doesn't shadow tool named 'search'. - - This tests a specific bug where doing 'import re' could cause - the 'search' tool to be replaced by re.search somehow. - """ - search_calls = [] - - def search_tool(query: str) -> str: - search_calls.append(query) - return f"results for: {query}" - - interpreter = JspiBackend( - preinstall_packages=False, - tools={"search": search_tool}, - ) - try: - # Use search tool before import - output1 = interpreter.execute(""" -r = await search("test query 1") -print(f"Result: {r}") -""") - assert "results for: test query 1" in str(output1) - assert len(search_calls) == 1 - - # Import re module (which has re.search) - output2 = interpreter.execute(""" -import re -pattern = re.compile(r"\\d+") -match = pattern.search("abc123def") -print(f"Found: {match.group()}") -""") - assert "Found: 123" in str(output2) - - # Search tool should still work and NOT be re.search - output3 = interpreter.execute(""" -r = await search("test query 2") -print(f"Result: {r}") -print(f"search is callable: {callable(search)}") -# Verify it's our async tool, not re.search -import inspect -is_coroutine = inspect.iscoroutinefunction(search) -print(f"search is async: {is_coroutine}") -""") - assert "results for: test query 2" in str(output3) - assert "search is callable: True" in str(output3) - assert "search is async: True" in str(output3) - assert len(search_calls) == 2 - - finally: - interpreter.shutdown() - - def test_tools_survive_multiple_consecutive_errors(self): - """Tools survive through multiple consecutive errors.""" - tool_calls = [] - - def my_tool(x: str) -> str: - tool_calls.append(x) - return f"result: {x}" - - interpreter = JspiBackend( - preinstall_packages=False, - tools={"my_tool": my_tool}, - ) - try: - # Error 1 - output1 = interpreter.execute(""" -try: - 1/0 -except ZeroDivisionError: - print("Error 1") -""") - assert "Error 1" in str(output1) - - # Tool still works - output2 = interpreter.execute(""" -r = await my_tool("after error 1") -print(r) -""") - assert "result: after error 1" in str(output2) - - # Error 2 - output3 = interpreter.execute(""" -try: - raise KeyError("missing") -except KeyError: - print("Error 2") -""") - assert "Error 2" in str(output3) - - # Tool still works - output4 = interpreter.execute(""" -r = await my_tool("after error 2") -print(r) -""") - assert "result: after error 2" in str(output4) - - # Error 3 - name error - output5 = interpreter.execute(""" -try: - undefined_variable_xyz -except NameError: - print("Error 3") -""") - assert "Error 3" in str(output5) + due_date: str | None = Field(default=None) + active: bool = True - # Tool still works - output6 = interpreter.execute(""" -r = await my_tool("after error 3") -print(r) -""") - assert "result: after error 3" in str(output6) - assert len(tool_calls) == 3 + items = [ + ExtractedItem(title="Task A", priority=None, due_date=None, active=True), + ExtractedItem(title="Task B", priority="high", due_date="2025-01-01", active=False), + ] + interpreter = JspiBackend(preinstall_packages=False) + try: + result = interpreter.execute( + """ +print(f"count={len(items)}") +print(f"a_priority={items[0]['priority']}") +print(f"a_due={items[0]['due_date']}") +print(f"a_active={items[0]['active']}") +print(f"b_priority={items[1]['priority']}") +print(f"b_active={items[1]['active']}") +""", + variables={"items": items}, + ) + output = str(result) + assert "count=2" in output + assert "a_priority=None" in output + assert "a_due=None" in output + assert "a_active=True" in output + assert "b_priority=high" in output + assert "b_active=False" in output finally: interpreter.shutdown() -class TestCancellationAndLateResponses: - """Tests for the resilient tool call channel with cancellation.""" - - def test_gather_with_one_failure_cancels_pending_calls(self): - """When one tool in asyncio.gather fails, pending calls are cancelled.""" - import asyncio - - tool_calls = [] - call_order = [] +class TestCustomPydanticTypesInSignatures: + def test_nested_pydantic_types_in_signature(self): + """Nested Pydantic models extract schemas with $defs for nested types.""" + received_schemas = [] - async def slow_tool(msg): - tool_calls.append(msg) - call_order.append(f"start:{msg}") - # Simulate slow tool - await asyncio.sleep(0.2) - if msg == "fail": - raise ValueError(f"Intentional failure for {msg}") - call_order.append(f"end:{msg}") - return f"result: {msg}" + def mock_predict(signature: str, pydantic_schemas=None, **kwargs): + received_schemas.append(pydantic_schemas) + return { + "person": {"name": "Alice", "address": {"street": "123 Main", "city": "NYC"}} + } - interpreter = JspiBackend( - preinstall_packages=False, - tools={"slow_tool": slow_tool}, - ) + interpreter = JspiBackend(tools={"predict": mock_predict}) try: - # This gather will fail because one tool raises an exception - # The error is wrapped by the tool call wrapper, so we catch RuntimeError - output = interpreter.execute(""" -import asyncio + result = interpreter.execute(""" +from pydantic import BaseModel -async def run(): - try: - results = await asyncio.gather( - slow_tool("call1"), - slow_tool("fail"), # This will fail - slow_tool("call3"), - ) - return results - except RuntimeError as e: - return f"Caught: {e}" +class Address(BaseModel): + street: str + city: str -result = await run() -print(result) -""") - # The error should be caught (wrapped as RuntimeError by tool wrapper) - assert "Caught:" in str(output) - assert "Intentional failure for fail" in str(output) - # All three calls were made - assert len(tool_calls) == 3 +class Person(BaseModel): + name: str + address: Address # Nested model +result = await predict("text: str -> person: Person", text="test") +print(f"Got person: {result['person'].name} at {result['person'].address.city}") +""") + assert "Got person: Alice at NYC" in str(result) + # Verify schema was extracted + assert len(received_schemas) == 1 + schemas = received_schemas[0] + assert schemas is not None + assert "Person" in schemas + # Verify nested type is in $defs + assert "$defs" in schemas["Person"] + assert "Address" in schemas["Person"]["$defs"] finally: interpreter.shutdown() - def test_next_execution_works_after_cancelled_gather(self): - """After a gather fails and cancels pending calls, next execution works.""" - import asyncio - call_count = [0] +class TestSerializationFailureRecovery: + def test_type_error_during_serialization(self): + """Tools survive TypeError during argument serialization.""" + received = [] - async def counting_tool(msg): - call_count[0] += 1 - await asyncio.sleep(0.05) - if msg == "fail": - raise ValueError("Intentional failure") - return f"result: {msg}" + def accept_dict(d: dict) -> str: + received.append(d) + return f"Got dict with {len(d)} keys" - interpreter = JspiBackend( - preinstall_packages=False, - tools={"counting_tool": counting_tool}, - ) + interpreter = JspiBackend(preinstall_packages=False, tools={"accept_dict": accept_dict}) try: - # First execution - gather with failure + # Success output1 = interpreter.execute(""" -import asyncio - -async def run(): - try: - results = await asyncio.gather( - counting_tool("a"), - counting_tool("fail"), - counting_tool("c"), - ) - return results - except RuntimeError as e: - return f"First caught: {e}" - -result = await run() +result = await accept_dict({"key": "value"}) print(result) """) - assert "First caught:" in str(output1) - assert "Intentional failure" in str(output1) - first_call_count = call_count[0] - assert first_call_count == 3 + assert "Got dict with 1 keys" in str(output1) - # Second execution - should work normally + # Try to pass something that looks like it has model_dump but fails output2 = interpreter.execute(""" -import asyncio +class FakeModel: + def model_dump(self): + raise TypeError("Cannot serialize") -async def run(): - results = await asyncio.gather( - counting_tool("x"), - counting_tool("y"), - ) - return results +try: + obj = FakeModel() + result = await accept_dict(obj) + print("Should not reach") +except Exception as e: + print(f"Caught: {type(e).__name__}") +""") + # Should catch the error (either TypeError or the wrapped error) + assert "Caught:" in str(output2) -result = await run() + # Tool should still work + output3 = interpreter.execute(""" +result = await accept_dict({"recovery": True}) print(result) """) - # Should succeed - assert "result: x" in str(output2) - assert "result: y" in str(output2) - # Two more calls made - assert call_count[0] == first_call_count + 2 - + assert "Got dict with 1 keys" in str(output3) finally: interpreter.shutdown() + +class TestCancellationAndLateResponses: def test_late_responses_are_ignored(self): """Late tool responses (after cancellation) are gracefully ignored.""" import asyncio @@ -2995,189 +386,17 @@ async def run(): finally: interpreter.shutdown() - def test_many_parallel_calls_one_fails(self): - """Test with 10 parallel calls where call 7 fails.""" - import asyncio - - call_log = [] - - async def numbered_tool(n): - call_log.append(f"start_{n}") - await asyncio.sleep(0.05) - if n == 7: - raise ValueError(f"Call {n} failed") - call_log.append(f"end_{n}") - return f"result_{n}" - - interpreter = JspiBackend( - preinstall_packages=False, - tools={"numbered_tool": numbered_tool}, - ) - try: - # Issue 10 parallel calls, call 7 will fail - output = interpreter.execute(""" -import asyncio - -async def run(): - try: - results = await asyncio.gather(*[ - numbered_tool(i) for i in range(1, 11) - ]) - return results - except RuntimeError as e: - return f"Failed: {e}" - -result = await run() -print(result) -""") - # Should catch the error from call 7 - assert "Failed:" in str(output) - assert "Call 7 failed" in str(output) - # All 10 calls were started - assert len([c for c in call_log if c.startswith("start_")]) == 10 - - # Next execution should work - output2 = interpreter.execute(""" -result = await numbered_tool(99) -print(result) -""") - assert "result_99" in str(output2) - - finally: - interpreter.shutdown() - - def test_rapid_failure_and_recovery(self): - """Test multiple rapid failures followed by successful execution.""" - import asyncio - - async def flaky_tool(should_fail): - await asyncio.sleep(0.02) - if should_fail: - raise RuntimeError("Flaky failure") - return "success" - - interpreter = JspiBackend( - preinstall_packages=False, - tools={"flaky_tool": flaky_tool}, - ) - try: - # Rapid failures - for i in range(3): - output = interpreter.execute( - f""" -import asyncio - -async def run(): - try: - await asyncio.gather( - flaky_tool(False), - flaky_tool(True), # Always fails - flaky_tool(False), - ) - except RuntimeError as e: - return f"Failure {i}: {{e}}" - -result = await run() -print(result) -""".replace("{i}", str(i)) - ) - # Error is wrapped by tool wrapper, check for the core message - assert f"Failure {i}:" in str(output) - assert "Flaky failure" in str(output) - - # Now success - output = interpreter.execute(""" -import asyncio - -results = await asyncio.gather( - flaky_tool(False), - flaky_tool(False), -) -print(results) -""") - assert "success" in str(output) - - finally: - interpreter.shutdown() - class TestPydanticReconstruction: - """Predict results should be reconstructed as Pydantic model instances.""" - - def test_predict_returns_pydantic_instances_for_list_output(self): - """When predict returns list[TaskItem], items should have attribute access.""" - - def mock_predict(signature: str, pydantic_schemas=None, **kwargs): - return {"tasks": [ - {"category": "Cert", "title": "Get ISO cert"}, - {"category": "Form", "title": "Fill W-9"}, - ]} - - interpreter = JspiBackend(tools={"predict": mock_predict}) - try: - result = interpreter.execute(""" -from pydantic import BaseModel - -class TaskItem(BaseModel): - category: str - title: str - -result = await predict("doc: str -> tasks: list[TaskItem]", doc="test") -# Attribute access should work — not just dict access -for task in result["tasks"]: - print(f"{task.title} ({task.category})") -""") - assert "Get ISO cert (Cert)" in str(result) - assert "Fill W-9 (Form)" in str(result) - finally: - interpreter.shutdown() - - def test_predict_returns_pydantic_instance_for_single_output(self): - """When predict returns a single Pydantic type, it should have attribute access.""" - - def mock_predict(signature: str, pydantic_schemas=None, **kwargs): - return {"summary": {"title": "Project X", "score": 95}} - - interpreter = JspiBackend(tools={"predict": mock_predict}) - try: - result = interpreter.execute(""" -from pydantic import BaseModel - -class Summary(BaseModel): - title: str - score: int - -result = await predict("doc: str -> summary: Summary", doc="test") -print(f"{result['summary'].title}: {result['summary'].score}") -""") - assert "Project X: 95" in str(result) - finally: - interpreter.shutdown() - - def test_predict_dict_output_stays_as_dict(self): - """When output type is plain dict (no Pydantic model), result stays as dict.""" - - def mock_predict(signature: str, pydantic_schemas=None, **kwargs): - return {"items": [{"a": 1}, {"b": 2}]} - - interpreter = JspiBackend(tools={"predict": mock_predict}) - try: - result = interpreter.execute(""" -result = await predict("doc: str -> items: list[dict]", doc="test") -# dict access should work -print(f"Got {len(result['items'])} items, first has key: {list(result['items'][0].keys())[0]}") -""") - assert "Got 2 items" in str(result) - finally: - interpreter.shutdown() - def test_can_add_fields_after_reconstruction(self): """LM can add metadata fields to reconstructed models (extra='allow').""" def mock_predict(signature: str, pydantic_schemas=None, **kwargs): - return {"tasks": [ - {"category": "Cert", "title": "Get ISO cert", "extra_field": "bonus"}, - ]} + return { + "tasks": [ + {"category": "Cert", "title": "Get ISO cert", "extra_field": "bonus"}, + ] + } interpreter = JspiBackend(tools={"predict": mock_predict}) try: @@ -3196,23 +415,6 @@ class TaskItem(BaseModel): finally: interpreter.shutdown() - def test_predict_result_attribute_access(self): - """Top-level result supports both attribute and subscript access.""" - - def mock_predict(signature: str, pydantic_schemas=None, **kwargs): - return {"answer": "Paris", "confidence": 0.95} - - interpreter = JspiBackend(tools={"predict": mock_predict}) - try: - result = interpreter.execute(""" -result = await predict("question: str -> answer: str, confidence: float", question="capital?") -# Both access patterns should work -print(f"attr: {result.answer}, sub: {result['confidence']}") -""") - assert "attr: Paris, sub: 0.95" in str(result) - finally: - interpreter.shutdown() - def test_predict_result_items_no_collision(self): """result.items returns stored list, not dict.items() method.""" @@ -3236,15 +438,6 @@ def mock_predict(signature: str, pydantic_schemas=None, **kwargs): class TestSandboxFatalErrors: - """Sandbox-level failures (exec timeout, BrokenPipe) must raise - SandboxFatalError, which is NOT a subclass of CodeInterpreterError. - - This ensures DSPy's RLM._execute_iteration catch on - (CodeInterpreterError, SyntaxError) cannot swallow them. The run - aborts instead of limping on with a freshly-spawned sandbox that - no longer has the per-run file_plan mounts. - """ - def test_exec_timeout_raises_sandbox_fatal_error(self): """exec_timeout firing raises SandboxFatalError, not CodeInterpreterError.""" interpreter = JspiBackend(preinstall_packages=False, exec_timeout=2.0) @@ -3254,22 +447,68 @@ def test_exec_timeout_raises_sandbox_fatal_error(self): finally: interpreter.shutdown() - def test_timeout_error_survives_rlm_catch_tuple(self): - """Simulate DSPy's RLM._execute_iteration handler: the fatal error - must propagate past `except (CodeInterpreterError, SyntaxError)`.""" - interpreter = JspiBackend(preinstall_packages=False, exec_timeout=2.0) - caught_by_rlm_handler = False - propagated = False - try: - try: - interpreter.execute("while True:\n pass\n") - except (CodeInterpreterError, SyntaxError): - caught_by_rlm_handler = True - except SandboxFatalError: - propagated = True - finally: - interpreter.shutdown() - assert propagated, "SandboxFatalError should propagate past RLM's handler" - assert not caught_by_rlm_handler, ( - "SandboxFatalError must not be caught by (CodeInterpreterError, SyntaxError)" + +class WorkspaceCancellationSignature(dspy.Signature): + workspace: Workspace = dspy.InputField() + answer: str = dspy.OutputField() + + +@pytest.mark.integration +@pytest.mark.skipif(shutil.which("deno") is None, reason="JSPI lifecycle test requires Deno") +@pytest.mark.asyncio +async def test_jspi_cancellation_flushes_mirror_before_sandbox_shutdown(tmp_path: Path): + mutation_completed = asyncio.Event() + + async def signal_mutation() -> str: + """Tell the host that the sandbox mutation completed.""" + asyncio.get_running_loop().call_later(0.1, mutation_completed.set) + return "ok" + + workspace_root = tmp_path / "workspace" + workspace_root.mkdir() + source = workspace_root / "source.txt" + source.write_text("before", encoding="utf-8") + + rlm = PredictRLM( + WorkspaceCancellationSignature, + lm=MagicMock(history=[]), + tools={"signal_mutation": signal_mutation}, + max_iterations=1, + verbose=False, + ) + rlm.generate_action.acall = AsyncMock( + return_value=dspy.Prediction( + reasoning="mutate the workspace and remain active", + code=( + "from pathlib import Path\n" + "Path('/sandbox/workspace/source.txt').write_text('after-cancel')\n" + "await signal_mutation()\n" + "while True:\n" + " pass" + ), ) + ) + rlm._configure_run_predictors = MagicMock() + + invocation = asyncio.create_task( + rlm.aforward(workspace=Workspace(path=str(workspace_root))) + ) + mutation_wait = asyncio.create_task(mutation_completed.wait()) + done, _ = await asyncio.wait( + {invocation, mutation_wait}, + timeout=30, + return_when=asyncio.FIRST_COMPLETED, + ) + if invocation in done: + await invocation + if mutation_wait not in done: + invocation.cancel() + mutation_wait.cancel() + await asyncio.gather(invocation, mutation_wait, return_exceptions=True) + pytest.fail("sandbox mutation did not complete before the test timeout") + invocation.cancel() + + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(invocation, timeout=30) + + assert source.read_text(encoding="utf-8") == "after-cancel" diff --git a/tests/test_interpreter_io.py b/tests/test_interpreter_io.py index 01e33217..78a7c61d 100644 --- a/tests/test_interpreter_io.py +++ b/tests/test_interpreter_io.py @@ -1,941 +1,104 @@ import asyncio -import errno -import json import os -import subprocess import threading -import time -import types +from types import SimpleNamespace import pytest from dspy.primitives.code_interpreter import CodeInterpreterError -import predict_rlm.backends.jspi.backend as rlm_interpreter from predict_rlm.backends import JspiBackend -from predict_rlm.backends.jspi.backend import JSONRPC_APP_ERRORS +from predict_rlm.backends.base import SandboxFatalError -class _BlockingStdin: - def write(self, data): - raise BlockingIOError(errno.EAGAIN, "pipe full") - - def flush(self): - raise AssertionError("flush should not run") - - -class _SilentStderr: - def read(self): - return "" - - -class _BufferingStdin: - def __init__(self): - self.data = [] - self.flushed = False - - def write(self, data): - self.data.append(data) - - def flush(self): - self.flushed = True - - -class _BufferingStdout: - def __init__(self, lines: list[str], fd: int = 123): - self.lines = list(lines) - self._fd = fd - - def fileno(self): - return self._fd - - def readline(self): - if self.lines: - return self.lines.pop(0) - return "" - - -def test_execute_recovers_when_stdin_blocks(monkeypatch): - interpreter = JspiBackend(preinstall_packages=False) - +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) +async def test_pipe_backpressure_does_not_drop_bytes(asynchronous): read_fd, write_fd = os.pipe() + os.set_blocking(write_fd, False) + backend = JspiBackend.__new__(JspiBackend) + backend._stdin_fd = write_fd + backend.deno_process = SimpleNamespace(poll=lambda: None, stdin=object()) + payload = "é" * 100_000 + "\n" + received = bytearray() + + def drain(): + while chunk := os.read(read_fd, 4096): + received.extend(chunk) + + reader = threading.Thread(target=drain) + reader.start() try: - interpreter.deno_process = types.SimpleNamespace( - stdin=_BlockingStdin(), - stderr=_SilentStderr(), - poll=lambda: None, - ) - interpreter._stdin_fd = write_fd - interpreter._stdout_fd = read_fd - interpreter._ensure_deno_process = lambda: None - interpreter._mount_files = lambda: None - interpreter._register_tools = lambda: None - interpreter._tools_registered = True - interpreter._mounted_files = True - - real_os_write = rlm_interpreter.os.write - state = {"blocked": True, "calls": 0} - - def fake_os_write(fd, data): - assert fd == write_fd - state["calls"] += 1 - if state["blocked"]: - state["blocked"] = False - raise BlockingIOError(errno.EAGAIN, "would block") - return real_os_write(fd, data) - - monkeypatch.setattr(rlm_interpreter.os, "write", fake_os_write) - - select_calls = {"count": 0} - - def fake_select(rlist, wlist, xlist, timeout=None): - select_calls["count"] += 1 - return ([], wlist, []) - - monkeypatch.setattr(rlm_interpreter.select, "select", fake_select) - - async def fake_execute_async(self, request_id): - return "ok" - - interpreter._execute_async = types.MethodType(fake_execute_async, interpreter) - - loop = asyncio.new_event_loop() - monkeypatch.setattr(asyncio, "get_event_loop", lambda: loop) - try: - result = interpreter.execute("print('hi')") - finally: - loop.close() - - assert result == "ok" - assert state["calls"] == 2 - assert select_calls["count"] == 1 + if asynchronous: + await asyncio.wait_for(backend._write_stdin_async(payload), timeout=5) + else: + await asyncio.wait_for(asyncio.to_thread(backend._write_stdin, payload), timeout=5) finally: - os.close(read_fd) os.close(write_fd) + reader.join(timeout=5) + os.close(read_fd) + assert not reader.is_alive() + assert received.decode() == payload -def test_send_request_falls_back_before_fd_ready(monkeypatch): - interpreter = JspiBackend(preinstall_packages=False) - stdin = _BufferingStdin() - - interpreter.deno_process = types.SimpleNamespace( - stdin=stdin, - stderr=None, - poll=lambda: None, - ) - interpreter._stdin_fd = -1 - - def fake_read(timeout=None): - return json.dumps({"id": interpreter._request_id, "result": {"output": "ok"}}) - - interpreter._read_with_timeout = fake_read # type: ignore[assignment] - - response = interpreter._send_request("execute", {"code": "print(1)"}, "during test") - - assert stdin.data # ensure blocking fallback wrote data - assert stdin.flushed - assert response["result"]["output"] == "ok" - - -def test_read_with_timeout_falls_back(monkeypatch): - interpreter = JspiBackend(preinstall_packages=False) - expected_line = json.dumps({"result": {"output": "hello"}, "id": 1}) + "\n" - stdout = _BufferingStdout([expected_line]) - - interpreter.deno_process = types.SimpleNamespace( - stdin=None, - stdout=stdout, - poll=lambda: None, - ) - interpreter._stdout_fd = -1 - interpreter._request_id = 1 - - def fake_select(rlist, wlist, xlist, timeout=None): - return (rlist, [], []) - - monkeypatch.setattr(rlm_interpreter.select, "select", fake_select) - - line = interpreter._read_with_timeout(timeout=0.1) - assert line == expected_line.strip() - - -# --------------------------------------------------------------------------- -# _get_semaphore classmethod -# --------------------------------------------------------------------------- - - -def test_get_semaphore_creates_on_first_call(): - """First call creates an asyncio.Semaphore; second call returns the same one.""" - # Reset class-level state so the test is isolated - JspiBackend._sandbox_semaphore = None - try: - sem1 = JspiBackend._get_semaphore() - sem2 = JspiBackend._get_semaphore() - assert isinstance(sem1, asyncio.Semaphore) - assert sem1 is sem2 - finally: - JspiBackend._sandbox_semaphore = None - - -# --------------------------------------------------------------------------- -# shutdown() with timeout -# --------------------------------------------------------------------------- - - -def test_shutdown_clean_exit(): - """shutdown() sends shutdown message and waits; process exits cleanly.""" - interpreter = JspiBackend(preinstall_packages=False) - wait_calls = [] - close_called = {"value": False} - - class _ClosableBufferingStdin(_BufferingStdin): - def close(self): - close_called["value"] = True - - stdin = _ClosableBufferingStdin() - - interpreter.deno_process = types.SimpleNamespace( - stdin=stdin, - stderr=_SilentStderr(), - poll=lambda: None, # process is alive - wait=lambda timeout=None: wait_calls.append(timeout), - kill=lambda: (_ for _ in ()).throw(AssertionError("kill should not be called")), - ) - interpreter._stdin_fd = -1 # use blocking fallback for _write_stdin - - interpreter.shutdown() - - assert any('"method": "shutdown"' in d or '"method":"shutdown"' in d for d in stdin.data) - assert close_called["value"] - assert interpreter.deno_process is None - assert wait_calls == [5] - - -def test_shutdown_timeout_kills_process(): - """shutdown() kills process when wait() raises TimeoutExpired.""" - interpreter = JspiBackend(preinstall_packages=False) - stdin = _BufferingStdin() - killed = {"value": False} - final_wait_calls = [] - - def fake_wait(timeout=None): - if timeout is not None: - raise subprocess.TimeoutExpired("deno", timeout) - final_wait_calls.append(True) - - # stdin.close() needs to be a callable - close_called = {"value": False} - - class _ClosableStdin(_BufferingStdin): - def close(self): - close_called["value"] = True - - stdin = _ClosableStdin() - - interpreter.deno_process = types.SimpleNamespace( - stdin=stdin, - stderr=_SilentStderr(), - poll=lambda: None, - wait=fake_wait, - kill=lambda: killed.__setitem__("value", True), - ) - interpreter._stdin_fd = -1 - - interpreter.shutdown() - - assert killed["value"] is True - assert len(final_wait_calls) == 1 # wait() after kill, no timeout - assert interpreter.deno_process is None - - -def test_shutdown_noop_when_process_already_exited(): - """shutdown() is a no-op if deno_process is None.""" - interpreter = JspiBackend(preinstall_packages=False) - interpreter.deno_process = None - interpreter.shutdown() # should not raise - assert interpreter.deno_process is None - - -# --------------------------------------------------------------------------- -# _sync_files -# --------------------------------------------------------------------------- - - -def test_sync_files_sends_messages(): - """_sync_files sends a sync_file JSON-RPC message for each write path.""" - interpreter = JspiBackend(preinstall_packages=False) - stdin = _BufferingStdin() +def test_request_deadline_survives_partial_stdout_and_preserves_buffer(monkeypatch): + import predict_rlm.backends.jspi.backend as backend_module - interpreter.deno_process = types.SimpleNamespace( - stdin=stdin, + monkeypatch.setattr(backend_module, "DENO_REQUEST_TIMEOUT_SEC", 0.1) + read_fd, write_fd = os.pipe() + backend = JspiBackend.__new__(JspiBackend) + backend._stdout_fd = read_fd + backend._stdin_fd = -1 + backend._read_buf = "" + backend._request_id = 0 + backend.deno_process = SimpleNamespace( + stdin=SimpleNamespace(write=lambda data: None, flush=lambda: None), + stdout=SimpleNamespace(fileno=lambda: read_fd), stderr=None, poll=lambda: None, ) - interpreter._stdin_fd = -1 - interpreter.enable_write_paths = ["/host/data/output.csv", "/host/data/report.json"] - interpreter.sync_files = True - - interpreter._sync_files() - - assert len(stdin.data) == 2 - for i, path in enumerate(interpreter.enable_write_paths): - msg = json.loads(stdin.data[i].rstrip("\n")) - assert msg["method"] == "sync_file" - assert msg["params"]["host_path"] == str(path) - assert msg["params"]["virtual_path"] == f"/sandbox/{os.path.basename(path)}" - - -def test_sync_files_skips_when_disabled(): - """_sync_files is a no-op when sync_files is False.""" - interpreter = JspiBackend(preinstall_packages=False) - stdin = _BufferingStdin() - - interpreter.deno_process = types.SimpleNamespace( - stdin=stdin, stderr=None, poll=lambda: None - ) - interpreter._stdin_fd = -1 - interpreter.enable_write_paths = ["/host/data/output.csv"] - interpreter.sync_files = False - - interpreter._sync_files() - - assert stdin.data == [] - - -def test_sync_files_skips_when_no_write_paths(): - """_sync_files is a no-op when enable_write_paths is empty.""" - interpreter = JspiBackend(preinstall_packages=False) - stdin = _BufferingStdin() - - interpreter.deno_process = types.SimpleNamespace( - stdin=stdin, stderr=None, poll=lambda: None - ) - interpreter._stdin_fd = -1 - interpreter.enable_write_paths = [] - interpreter.sync_files = True - - interpreter._sync_files() - - assert stdin.data == [] - - -# --------------------------------------------------------------------------- -# _send_request error paths -# --------------------------------------------------------------------------- - - -def test_send_request_deno_exit_detection(): - """_send_request raises CodeInterpreterError when Deno has exited.""" - interpreter = JspiBackend(preinstall_packages=False) - stdin = _BufferingStdin() - - # poll() returns None during _write_stdin (process alive), then 1 after - # _read_with_timeout returns empty (process died mid-execution). - poll_results = iter([None, 1]) - - interpreter.deno_process = types.SimpleNamespace( - stdin=stdin, - stderr=types.SimpleNamespace(read=lambda: "segfault"), - poll=lambda: next(poll_results), - ) - interpreter._stdin_fd = -1 - - # _read_with_timeout returns empty string (Deno died) - interpreter._read_with_timeout = lambda timeout=None: "" # type: ignore[assignment] - - with pytest.raises(CodeInterpreterError, match=r"Deno exited \(code 1\).*segfault"): - interpreter._send_request("execute", {"code": "1+1"}, "during test") - + errors = [] -def test_send_request_response_id_mismatch(): - """_send_request raises CodeInterpreterError when no matching response - ever arrives. The resync loop discards mismatched-id frames (stale - responses from prior timed-out requests) and only raises after - exhausting its safety cap — ensuring that a runaway wrong-id stream - doesn't hang the caller forever. - """ - interpreter = JspiBackend(preinstall_packages=False) - stdin = _BufferingStdin() - - interpreter.deno_process = types.SimpleNamespace( - stdin=stdin, stderr=None, poll=lambda: None - ) - interpreter._stdin_fd = -1 - - # Return a response with mismatched id forever — the resync cap - # ensures we bail instead of spinning. - interpreter._read_with_timeout = lambda timeout=None: json.dumps( # type: ignore[assignment] - {"id": 9999, "result": {"output": "ok"}} - ) - - with pytest.raises(CodeInterpreterError, match="stale|resync"): - interpreter._send_request("execute", {"code": "1+1"}, "during test") - - -def test_send_request_error_in_response(): - """_send_request raises CodeInterpreterError when response contains an error.""" - interpreter = JspiBackend(preinstall_packages=False) - stdin = _BufferingStdin() - - interpreter.deno_process = types.SimpleNamespace( - stdin=stdin, stderr=None, poll=lambda: None - ) - interpreter._stdin_fd = -1 - - def fake_read(timeout=None): - return json.dumps( - {"id": interpreter._request_id, "error": {"message": "tool failed"}} - ) - - interpreter._read_with_timeout = fake_read # type: ignore[assignment] - - with pytest.raises(CodeInterpreterError, match="tool failed"): - interpreter._send_request("execute", {"code": "1+1"}, "during test") - - -# --------------------------------------------------------------------------- -# top-level execution gate -# --------------------------------------------------------------------------- - - -def test_execute_serializes_concurrent_calls_without_real_deno(monkeypatch): - interpreter = JspiBackend(preinstall_packages=False) - interpreter.deno_process = types.SimpleNamespace( - stdin=_BufferingStdin(), - stderr=_SilentStderr(), - poll=lambda: None, - ) - interpreter._stdin_fd = -1 - interpreter._ensure_deno_process = lambda: None - interpreter._mount_files = lambda: None - interpreter._register_tools = lambda: None - - active = 0 - max_active = 0 - active_lock = threading.Lock() - loops: list[asyncio.AbstractEventLoop] = [] - - async def fake_execute_with_timeout(request_id, execute_start_time=None): - nonlocal active, max_active - with active_lock: - active += 1 - max_active = max(max_active, active) + def request(): try: - await asyncio.sleep(0.1) - return f"result-{request_id}" - finally: - with active_lock: - active -= 1 - - def new_event_loop(): - loop = asyncio.new_event_loop() - loops.append(loop) - return loop - - interpreter._execute_with_timeout = fake_execute_with_timeout # type: ignore[assignment] - monkeypatch.setattr(asyncio, "get_event_loop", new_event_loop) - - barrier = threading.Barrier(3) - results: list[str] = [] - errors: list[BaseException] = [] - - def run_execute() -> None: - barrier.wait() - try: - results.append(interpreter.execute("print('hi')")) + backend._send_request("health_check", {}, context="partial stdout") except BaseException as exc: errors.append(exc) - threads = [threading.Thread(target=run_execute) for _ in range(2)] + os.write(write_fd, b"partial") + worker = threading.Thread(target=request, daemon=True) + worker.start() try: - for thread in threads: - thread.start() - barrier.wait() - for thread in threads: - thread.join(timeout=2) + worker.join(timeout=2) + assert not worker.is_alive(), "partial stdout bypassed the request deadline" + assert len(errors) == 1 and isinstance(errors[0], CodeInterpreterError) + os.write(write_fd, b"-completion\n") + assert backend._read_line_raw(timeout=1) == "partial-completion" finally: - for loop in loops: - loop.close() - - assert [thread.is_alive() for thread in threads] == [False, False] - assert errors == [] - assert sorted(results) == ["result-1", "result-2"] - assert max_active == 1 - - -@pytest.mark.asyncio -async def test_aexecute_serializes_concurrent_calls_without_real_deno(): - JspiBackend._sandbox_semaphore = None - try: - interpreter = JspiBackend(preinstall_packages=False) - active = 0 - max_active = 0 - - async def fake_inner(code, variables): - nonlocal active, max_active - active += 1 - max_active = max(max_active, active) - try: - await asyncio.sleep(0.1) - return code - finally: - active -= 1 - - interpreter._aexecute_inner = fake_inner # type: ignore[assignment] - - start = time.monotonic() - results = await asyncio.gather( - interpreter.aexecute("print('one')"), - interpreter.aexecute("print('two')"), - ) - - assert sorted(results) == ["print('one')", "print('two')"] - assert max_active == 1 - assert time.monotonic() - start >= 0.18 - finally: - JspiBackend._sandbox_semaphore = None - - -@pytest.mark.asyncio -async def test_aexecute_from_tool_callback_context_raises_runtimeerror(): - interpreter = JspiBackend(preinstall_packages=False) - - with interpreter._execution_gate.async_tool_callback(): - with pytest.raises(RuntimeError, match="host tool callback"): - await interpreter.aexecute("print('nested')") - - -# --------------------------------------------------------------------------- -# aexecute semaphore -# --------------------------------------------------------------------------- - - -@pytest.mark.asyncio -async def test_aexecute_acquires_and_releases_semaphore(): - """aexecute acquires the semaphore before running and releases it after.""" - JspiBackend._sandbox_semaphore = None - try: - interpreter = JspiBackend(preinstall_packages=False) - sem = JspiBackend._get_semaphore() - - acquire_count = {"value": 0} - release_count = {"value": 0} - original_acquire = sem.acquire - original_release = sem.release - - async def tracking_acquire(): - acquire_count["value"] += 1 - return await original_acquire() - - def tracking_release(): - release_count["value"] += 1 - return original_release() - - sem.acquire = tracking_acquire # type: ignore[assignment] - sem.release = tracking_release # type: ignore[assignment] - - # Mock _aexecute_inner to avoid real Deno subprocess - async def fake_inner(code, variables): - return "result" - - interpreter._aexecute_inner = fake_inner # type: ignore[assignment] - - result = await interpreter.aexecute("print(1)") - assert result == "result" - assert acquire_count["value"] == 1 - assert release_count["value"] == 1 - finally: - JspiBackend._sandbox_semaphore = None - - -@pytest.mark.asyncio -async def test_aexecute_releases_semaphore_on_error(): - """aexecute releases the semaphore even when _aexecute_inner raises.""" - JspiBackend._sandbox_semaphore = None - try: - interpreter = JspiBackend(preinstall_packages=False) - sem = JspiBackend._get_semaphore() - - release_count = {"value": 0} - original_release = sem.release - - def tracking_release(): - release_count["value"] += 1 - return original_release() - - sem.release = tracking_release # type: ignore[assignment] - - async def failing_inner(code, variables): - raise RuntimeError("boom") - - interpreter._aexecute_inner = failing_inner # type: ignore[assignment] - - with pytest.raises(RuntimeError, match="boom"): - await interpreter.aexecute("print(1)") - - assert release_count["value"] == 1 - finally: - JspiBackend._sandbox_semaphore = None - - -# --------------------------------------------------------------------------- -# _write_stdin_async -# --------------------------------------------------------------------------- - - -@pytest.mark.asyncio -async def test_write_stdin_async_normal_write(monkeypatch): - """_write_stdin_async writes data via os.write in a single pass.""" - interpreter = JspiBackend(preinstall_packages=False) - written_chunks = [] - - interpreter.deno_process = types.SimpleNamespace(poll=lambda: None) - interpreter._stdin_fd = 42 - - def fake_write(fd, data): - assert fd == 42 - written_chunks.append(data) - return len(data) - - monkeypatch.setattr(rlm_interpreter.os, "write", fake_write) - - await interpreter._write_stdin_async("hello\n") - - assert b"".join(written_chunks) == b"hello\n" - - -@pytest.mark.asyncio -async def test_write_stdin_async_blocking_retry(monkeypatch): - """_write_stdin_async retries via add_writer when BlockingIOError occurs.""" - interpreter = JspiBackend(preinstall_packages=False) - - interpreter.deno_process = types.SimpleNamespace(poll=lambda: None) - interpreter._stdin_fd = 42 - - state = {"blocked": True} - written_chunks = [] - - def fake_write(fd, data): - if state["blocked"]: - state["blocked"] = False - raise BlockingIOError(errno.EAGAIN, "pipe full") - written_chunks.append(data) - return len(data) - - monkeypatch.setattr(rlm_interpreter.os, "write", fake_write) - - # We need a real event loop with add_writer support. - # The default event loop on macOS uses kqueue which requires real fds, - # so we mock add_writer/remove_writer to immediately signal writable. - loop = asyncio.get_running_loop() - - def fake_add_writer(fd, callback, *args): - # Immediately schedule the callback so the write can proceed - loop.call_soon(callback, *args) - - def fake_remove_writer(fd): - pass - - monkeypatch.setattr(loop, "add_writer", fake_add_writer) - monkeypatch.setattr(loop, "remove_writer", fake_remove_writer) - - await interpreter._write_stdin_async("data\n") - - assert b"".join(written_chunks) == b"data\n" - - -@pytest.mark.asyncio -async def test_write_stdin_async_raises_when_process_dead(): - """_write_stdin_async raises CodeInterpreterError when process has exited.""" - interpreter = JspiBackend(preinstall_packages=False) - interpreter.deno_process = types.SimpleNamespace(poll=lambda: 1) - interpreter._stdin_fd = 42 - - with pytest.raises(CodeInterpreterError, match="no longer running"): - await interpreter._write_stdin_async("hello\n") - - -# --------------------------------------------------------------------------- -# _send_completed_responses -# --------------------------------------------------------------------------- - - -@pytest.mark.asyncio -async def test_send_completed_responses_sends_one_result(monkeypatch): - """_send_completed_responses sends at most one response for a completed task.""" - interpreter = JspiBackend(preinstall_packages=False) - sent_messages = [] - - async def fake_write(data): - sent_messages.append(data) - - interpreter._write_stdin_async = fake_write # type: ignore[assignment] - - # Create two completed tasks - task1 = asyncio.get_running_loop().create_future() - task1.set_result({"value": "result1", "type": "string"}) - task2 = asyncio.get_running_loop().create_future() - task2.set_result({"value": "result2", "type": "string"}) - - pending = {"req1": task1, "req2": task2} - - await interpreter._send_completed_responses(pending) - - # Only one sent (at most one per call) - assert len(sent_messages) == 1 - msg = json.loads(sent_messages[0].rstrip("\n")) - assert "result" in msg - # One task was popped - assert len(pending) == 1 - - -@pytest.mark.asyncio -async def test_send_completed_responses_error_task(monkeypatch): - """_send_completed_responses sends a JSON-RPC error for a failed task.""" - interpreter = JspiBackend(preinstall_packages=False) - sent_messages = [] - - async def fake_write(data): - sent_messages.append(data) - - interpreter._write_stdin_async = fake_write # type: ignore[assignment] - - # Task that returned an error dict - task = asyncio.get_running_loop().create_future() - task.set_result({"error": "tool exploded"}) - - pending = {"req1": task} - - await interpreter._send_completed_responses(pending) - - assert len(sent_messages) == 1 - msg = json.loads(sent_messages[0].rstrip("\n")) - assert "error" in msg - assert msg["error"]["message"] == "tool exploded" - assert msg["error"]["code"] == JSONRPC_APP_ERRORS["RuntimeError"] - assert len(pending) == 0 - - -@pytest.mark.asyncio -async def test_send_completed_responses_exception_in_task(): - """_send_completed_responses handles an exception raised by the task.""" - interpreter = JspiBackend(preinstall_packages=False) - sent_messages = [] - - async def fake_write(data): - sent_messages.append(data) - - interpreter._write_stdin_async = fake_write # type: ignore[assignment] - - task = asyncio.get_running_loop().create_future() - task.set_exception(ValueError("unexpected")) - - pending = {"req1": task} - - await interpreter._send_completed_responses(pending) - - assert len(sent_messages) == 1 - msg = json.loads(sent_messages[0].rstrip("\n")) - assert "error" in msg - assert "unexpected" in msg["error"]["message"] - - -# --------------------------------------------------------------------------- -# _wait_and_send_all_responses -# --------------------------------------------------------------------------- - - -@pytest.mark.asyncio -async def test_wait_and_send_all_responses_sends_all(): - """_wait_and_send_all_responses waits for all tasks and sends responses.""" - interpreter = JspiBackend(preinstall_packages=False) - sent_messages = [] - - async def fake_write(data): - sent_messages.append(data) - - interpreter._write_stdin_async = fake_write # type: ignore[assignment] - - task1 = asyncio.get_running_loop().create_future() - task1.set_result({"value": "r1", "type": "string"}) - task2 = asyncio.get_running_loop().create_future() - task2.set_result({"value": "r2", "type": "json"}) - - pending = {"req1": task1, "req2": task2} - - await interpreter._wait_and_send_all_responses(pending) - - assert len(sent_messages) == 2 - assert len(pending) == 0 # cleared - - -@pytest.mark.asyncio -async def test_wait_and_send_all_responses_error_path(): - """_wait_and_send_all_responses sends error for a task that raises.""" - interpreter = JspiBackend(preinstall_packages=False) - sent_messages = [] - - async def fake_write(data): - sent_messages.append(data) - - interpreter._write_stdin_async = fake_write # type: ignore[assignment] - - task = asyncio.get_running_loop().create_future() - task.set_exception(RuntimeError("crash")) - - pending = {"req1": task} - - await interpreter._wait_and_send_all_responses(pending) - - assert len(sent_messages) == 1 - msg = json.loads(sent_messages[0].rstrip("\n")) - assert "error" in msg - assert "crash" in msg["error"]["message"] - assert len(pending) == 0 - - -# --------------------------------------------------------------------------- -# SyntaxError detail formatting in _execute_async -# --------------------------------------------------------------------------- - - -@pytest.mark.asyncio -async def test_execute_async_syntax_error_formatting(): - """_execute_async formats SyntaxError with line/col/text from args tuple.""" - interpreter = JspiBackend(preinstall_packages=False) - - # Build a JSON-RPC error response with SyntaxError args format - error_response = json.dumps({ - "jsonrpc": "2.0", - "id": 1, - "error": { - "code": -32000, - "message": "invalid syntax", - "data": { - "type": "SyntaxError", - "args": [ - "invalid syntax", - ["", 5, 10, "x = (1 +\n"], - ], - }, - }, - }) - - call_count = {"value": 0} - - async def fake_read(timeout=None): - call_count["value"] += 1 - if call_count["value"] == 1: - return error_response - return None - - interpreter._read_with_timeout_async = fake_read # type: ignore[assignment] - - # _send_completed_responses and _wait_and_send_all_responses need to be no-ops - async def noop_responses(pending): - pass - - interpreter._send_completed_responses = noop_responses # type: ignore[assignment] - interpreter._wait_and_send_all_responses = noop_responses # type: ignore[assignment] - - with pytest.raises(SyntaxError) as exc_info: - await interpreter._execute_async(1) - - detail = str(exc_info.value) - assert "line 5" in detail - assert "col 10" in detail - assert "x = (1 +" in detail + os.close(write_fd) + worker.join(timeout=2) + os.close(read_fd) -@pytest.mark.asyncio -async def test_execute_async_syntax_error_minimal_args(): - """_execute_async handles SyntaxError with minimal args (just message).""" - interpreter = JspiBackend(preinstall_packages=False) +def test_stdout_eof_does_not_drain_live_stderr(): + backend = JspiBackend.__new__(JspiBackend) - error_response = json.dumps({ - "jsonrpc": "2.0", - "id": 1, - "error": { - "code": -32000, - "message": "unexpected EOF", - "data": { - "type": "SyntaxError", - "args": ["unexpected EOF"], - }, - }, - }) + class BlockingStderr: + def read(self): + raise AssertionError("stderr.read() would block while its writer remains alive") - call_count = {"value": 0} - - async def fake_read(timeout=None): - call_count["value"] += 1 - if call_count["value"] == 1: - return error_response + async def no_completed_responses(_pending_tasks): return None - interpreter._read_with_timeout_async = fake_read # type: ignore[assignment] - - async def noop_responses(pending): - pass - - interpreter._send_completed_responses = noop_responses # type: ignore[assignment] - interpreter._wait_and_send_all_responses = noop_responses # type: ignore[assignment] - - with pytest.raises(SyntaxError, match="unexpected EOF"): - await interpreter._execute_async(1) - - -# --------------------------------------------------------------------------- -# _read_with_timeout fallback when fd < 0 and stdout.fileno() is used -# --------------------------------------------------------------------------- - - -def test_read_with_timeout_uses_stdout_fileno_when_fd_negative(monkeypatch): - """_read_with_timeout uses stdout.fileno() when _stdout_fd < 0.""" - interpreter = JspiBackend(preinstall_packages=False) - expected_line = json.dumps({"result": "ok", "id": 1}) + "\n" - - # _BufferingStdout.fileno() returns 999 by default - stdout = _BufferingStdout([expected_line], fd=999) - - interpreter.deno_process = types.SimpleNamespace( - stdin=None, - stdout=stdout, - poll=lambda: None, - ) - interpreter._stdout_fd = -1 - interpreter._request_id = 1 - - fileno_calls = {"count": 0} - original_fileno = stdout.fileno - - def tracking_fileno(): - fileno_calls["count"] += 1 - return original_fileno() - - stdout.fileno = tracking_fileno - - def fake_select(rlist, wlist, xlist, timeout=None): - # Verify select was called with the fileno() result - assert rlist == [999] - return (rlist, [], []) - - monkeypatch.setattr(rlm_interpreter.select, "select", fake_select) - - line = interpreter._read_with_timeout(timeout=0.1) - assert line == expected_line.strip() - assert fileno_calls["count"] == 1 - - -def test_read_with_timeout_returns_none_when_no_stdout(monkeypatch): - """_read_with_timeout returns None when stdout is None and fd < 0.""" - interpreter = JspiBackend(preinstall_packages=False) + async def stdout_eof(_timeout): + return "" - interpreter.deno_process = types.SimpleNamespace( - stdin=None, - stdout=None, - poll=lambda: None, + backend.deno_process = SimpleNamespace( + stderr=BlockingStderr(), + kill=lambda: None, + poll=lambda: 0, ) - interpreter._stdout_fd = -1 - - result = interpreter._read_with_timeout(timeout=0.1) - assert result is None + backend._pending_file_ops = {} + backend._send_completed_responses = no_completed_responses + backend._read_with_timeout_async = stdout_eof + with pytest.raises(SandboxFatalError): + asyncio.run(backend._execute_async(execute_request_id=1)) diff --git a/tests/test_interpreter_unit.py b/tests/test_interpreter_unit.py deleted file mode 100644 index 43384e0e..00000000 --- a/tests/test_interpreter_unit.py +++ /dev/null @@ -1,345 +0,0 @@ -"""Unit tests for interpreter helpers (no Deno required).""" - -import asyncio -import subprocess -import tempfile -import types -from types import SimpleNamespace -from typing import Any -from unittest.mock import patch - -import pytest -from dspy.primitives.code_interpreter import CodeInterpreterError - -from predict_rlm.backends import JspiBackend -from predict_rlm.backends.base import SandboxFatalError -from predict_rlm.backends.jspi.backend import RUNNER_PATH, _needs_jspi_flag -from predict_rlm.telemetry import TelemetryContext - - -class ListTelemetrySink: - def __init__(self): - self.records: list[dict[str, Any]] = [] - - def write(self, record: dict[str, Any]) -> None: - self.records.append(record) - - -class TestNeedsJspiFlag: - @patch.object(subprocess, "check_output") - def test_old_v8_needs_flag(self, mock_check): - mock_check.return_value = "deno 2.0.0\nv8 12.9.245.12-rusty\ntypescript 5.6.2" - assert _needs_jspi_flag() is True - - @patch.object(subprocess, "check_output") - def test_v8_13_6_needs_flag(self, mock_check): - mock_check.return_value = "deno 2.1.0\nv8 13.6.100.0\ntypescript 5.6.2" - assert _needs_jspi_flag() is True - - @patch.object(subprocess, "check_output") - def test_v8_13_7_no_flag(self, mock_check): - mock_check.return_value = "deno 2.2.0\nv8 13.7.0.0\ntypescript 5.6.2" - assert _needs_jspi_flag() is False - - @patch.object(subprocess, "check_output") - def test_v8_14_0_no_flag(self, mock_check): - mock_check.return_value = "deno 3.0.0\nv8 14.0.0.0\ntypescript 5.6.2" - assert _needs_jspi_flag() is False - - @patch.object(subprocess, "check_output") - def test_deno_not_found_returns_true(self, mock_check): - mock_check.side_effect = FileNotFoundError("deno not found") - assert _needs_jspi_flag() is True - - @patch.object(subprocess, "check_output") - def test_unexpected_output_returns_true(self, mock_check): - mock_check.return_value = "some garbage output" - assert _needs_jspi_flag() is True - - -def _make_interpreter(): - """Create a JspiBackend without running __init__ (no Deno subprocess).""" - return JspiBackend.__new__(JspiBackend) - - -def _attach_telemetry(interp: JspiBackend) -> ListTelemetrySink: - sink = ListTelemetrySink() - interp._telemetry_context = TelemetryContext(sink=sink, trace_id="trace-1") - interp._interpreter_id = "jspi-test" - return sink - - -class TestBuildDenoCommand: - @patch("predict_rlm.backends.jspi.backend._needs_jspi_flag", return_value=True) - def test_includes_jspi_flag_when_needed(self, _): - interp = _make_interpreter() - with patch.object(interp, "_get_deno_dir", return_value=[]): - cmd = interp._build_deno_command([], [], [], []) - assert "--v8-flags=--experimental-wasm-jspi" in cmd - - @patch("predict_rlm.backends.jspi.backend._needs_jspi_flag", return_value=False) - def test_excludes_jspi_flag_when_not_needed(self, _): - interp = _make_interpreter() - with patch.object(interp, "_get_deno_dir", return_value=[]): - cmd = interp._build_deno_command([], [], [], []) - assert "--v8-flags=--experimental-wasm-jspi" not in cmd - - @patch("predict_rlm.backends.jspi.backend._needs_jspi_flag", return_value=False) - def test_runner_path_in_command(self, _): - interp = _make_interpreter() - with patch.object(interp, "_get_deno_dir", return_value=[]): - cmd = interp._build_deno_command([], [], [], []) - assert str(RUNNER_PATH) in cmd - - @patch("predict_rlm.backends.jspi.backend._needs_jspi_flag", return_value=False) - def test_allow_read_includes_runner_and_user_paths(self, _): - interp = _make_interpreter() - with patch.object(interp, "_get_deno_dir", return_value=[]): - cmd = interp._build_deno_command(["/data/input"], [], [], []) - read_arg = [a for a in cmd if a.startswith("--allow-read=")][0] - read_paths = read_arg.split("=", 1)[1].split(",") - assert str(RUNNER_PATH) in read_paths - assert "/data/input" in read_paths - - @patch("predict_rlm.backends.jspi.backend._needs_jspi_flag", return_value=False) - def test_write_paths_also_in_allow_read(self, _): - interp = _make_interpreter() - with patch.object(interp, "_get_deno_dir", return_value=[]): - cmd = interp._build_deno_command([], ["/data/output"], [], []) - read_arg = [a for a in cmd if a.startswith("--allow-read=")][0] - assert "/data/output" in read_arg - - @patch("predict_rlm.backends.jspi.backend._needs_jspi_flag", return_value=False) - def test_allow_write_includes_tempdir_and_user_paths(self, _): - interp = _make_interpreter() - with patch.object(interp, "_get_deno_dir", return_value=[]): - cmd = interp._build_deno_command([], ["/data/output"], [], []) - write_arg = [a for a in cmd if a.startswith("--allow-write=")][0] - write_paths = write_arg.split("=", 1)[1].split(",") - assert "/data/output" in write_paths - assert tempfile.gettempdir() in write_paths - assert "/tmp" in write_paths - - @patch("predict_rlm.backends.jspi.backend._needs_jspi_flag", return_value=False) - def test_allow_net_with_domains(self, _): - interp = _make_interpreter() - with patch.object(interp, "_get_deno_dir", return_value=[]): - cmd = interp._build_deno_command([], [], ["pypi.org", "api.example.com"], []) - assert "--allow-net=pypi.org,api.example.com" in cmd - - @patch("predict_rlm.backends.jspi.backend._needs_jspi_flag", return_value=False) - def test_no_allow_net_when_empty(self, _): - interp = _make_interpreter() - with patch.object(interp, "_get_deno_dir", return_value=[]): - cmd = interp._build_deno_command([], [], [], []) - assert not any(a.startswith("--allow-net") for a in cmd) - - @patch("predict_rlm.backends.jspi.backend._needs_jspi_flag", return_value=False) - def test_always_includes_allow_env_and_no_prompt(self, _): - interp = _make_interpreter() - with patch.object(interp, "_get_deno_dir", return_value=[]): - cmd = interp._build_deno_command([], [], [], []) - assert "--allow-env" in cmd - assert "--no-prompt" in cmd - - @patch("predict_rlm.backends.jspi.backend._needs_jspi_flag", return_value=False) - def test_env_vars_as_final_arg(self, _): - interp = _make_interpreter() - with patch.object(interp, "_get_deno_dir", return_value=[]): - cmd = interp._build_deno_command( - [], [], [], ["PYODIDE_PREINSTALL", "SKILL_PACKAGES"] - ) - assert cmd[-1] == "PYODIDE_PREINSTALL,SKILL_PACKAGES" - - @patch("predict_rlm.backends.jspi.backend._needs_jspi_flag", return_value=False) - def test_no_env_vars_runner_is_last(self, _): - interp = _make_interpreter() - with patch.object(interp, "_get_deno_dir", return_value=[]): - cmd = interp._build_deno_command([], [], [], []) - assert cmd[-1] == str(RUNNER_PATH) - - -class TestGetDenoDir: - def test_includes_home_cache_paths(self): - interp = _make_interpreter() - with patch.dict("os.environ", {"HOME": "/home/test"}, clear=False): - dirs = interp._get_deno_dir() - assert "/home/test/.cache/deno" in dirs - assert "/home/test/Library/Caches/deno" in dirs - - def test_includes_deno_dir_env(self): - interp = _make_interpreter() - with patch.dict("os.environ", {"DENO_DIR": "/custom/deno"}, clear=False): - dirs = interp._get_deno_dir() - assert "/custom/deno" in dirs - - -class TestSandboxFatalError: - """SandboxFatalError must NOT inherit from CodeInterpreterError. - - DSPy's RLM._execute_iteration catches (CodeInterpreterError, SyntaxError) - and converts the exception into an "[Error] ..." string that gets fed - back to the model as regular iteration output. For ordinary in-sandbox - errors (NameError, tool raised, etc.) that's correct — the model can - self-correct. But when the sandbox subprocess itself dies (exec timeout, - BrokenPipe), the per-run file_plan mounts and output dirs are gone, so - subsequent iterations trip over FileNotFoundError with no way to recover. - - Keeping SandboxFatalError a sibling of CodeInterpreterError ensures the - base class's catch tuple does not swallow it — it propagates out of - rlm.forward() and the run fails fast. - """ - - def test_is_runtime_error(self): - assert issubclass(SandboxFatalError, RuntimeError) - - def test_is_not_code_interpreter_error(self): - assert not issubclass(SandboxFatalError, CodeInterpreterError) - - -class TestJspiLoggingConfig: - def test_configure_debug_and_verbose_are_independent(self): - interp = _make_interpreter() - interp._debug = False - interp._verbose = False - - interp.configure_debug(True) - assert interp._debug is True - assert interp._verbose is False - - interp.configure_verbose(True) - assert interp._debug is True - assert interp._verbose is True - - interp.configure_debug(False) - assert interp._debug is False - assert interp._verbose is True - - def test_configure_runtime_updates_debug_and_verbose(self): - interp = _make_interpreter() - interp._debug = False - interp._verbose = False - - interp.configure_runtime(debug=True, verbose=True) - - assert interp._debug is True - assert interp._verbose is True - - -class TestJspiTelemetry: - def test_health_check_no_response_emits_lifecycle_failure(self): - interp = _make_interpreter() - sink = _attach_telemetry(interp) - interp._request_id = 0 - interp._stdin_fd = -1 - interp._stdout_fd = -1 - interp._read_buf = "" - interp.deno_process = types.SimpleNamespace( - stdin=types.SimpleNamespace(write=lambda _data: None, flush=lambda: None), - stdout=None, - poll=lambda: None, - pid=1234, - ) - - with patch.object(interp, "_read_with_timeout", return_value=None): - with pytest.raises(CodeInterpreterError, match="No response"): - interp._send_request("health_check", {}, "during health check") - - names = [record["name"] for record in sink.records] - assert names == [ - "sandbox.health_check.start", - "sandbox.health_check.no_response", - ] - no_response = sink.records[-1] - assert no_response["status"]["code"] == "ERROR" - assert no_response["attributes"]["failure.class"] == "sandbox_lifecycle_failure" - assert no_response["attributes"]["process.pid"] == 1234 - assert no_response["attributes"]["rpc.request_id"] == 1 - - def test_execute_timeout_emits_timeout_and_kill_events(self): - interp = _make_interpreter() - sink = _attach_telemetry(interp) - interp._exec_timeout = 0.01 - interp._pending_file_ops = {} - interp.deno_process = types.SimpleNamespace( - kill=lambda: None, - poll=lambda: 0, - wait=lambda timeout=None: None, - pid=4321, - ) - - async def _never_returns(_request_id): - await asyncio.sleep(60) - - interp._execute_async = _never_returns - - with pytest.raises(SandboxFatalError): - asyncio.run(interp._execute_with_timeout(7, 1_000_000_000)) - - names = [record["name"] for record in sink.records] - assert "sandbox.execute.timeout" in names - assert "sandbox.shutdown.kill" in names - timeout = next( - record for record in sink.records if record["name"] == "sandbox.execute.timeout" - ) - assert timeout["status"]["code"] == "ERROR" - assert timeout["attributes"]["failure.class"] == "sandbox_exec_timeout" - assert timeout["attributes"]["rpc.request_id"] == 7 - assert timeout["attributes"]["process.pid"] == 4321 - - def test_tool_timeout_emits_host_tool_failure(self, monkeypatch): - import predict_rlm.backends.jspi.backend as rlm_interpreter - - monkeypatch.setattr(rlm_interpreter, "TOOL_CALL_TIMEOUT_SEC", 0.01) - - async def slow_tool(): - await asyncio.sleep(60) - - interp = _make_interpreter() - sink = _attach_telemetry(interp) - interp.tools = {"slow_tool": slow_tool} - interp._debug = False - interp._pending_file_ops = {} - - response = asyncio.run( - interp._execute_tool_async("slow_tool", {"args": [], "kwargs": {}}, "tool-1") - ) - - assert "error" in response - names = [record["name"] for record in sink.records] - assert names == ["sandbox.tool_call.start", "sandbox.tool_call.timeout"] - timeout = sink.records[-1] - assert timeout["attributes"]["failure.class"] == "host_tool_timeout_or_leak" - assert timeout["attributes"]["tool.name"] == "slow_tool" - assert timeout["attributes"]["tool.id"] == "tool-1" - - -class TestAsyncExecuteEof: - def test_stdout_eof_does_not_drain_stderr_with_unbounded_read(self): - interp = _make_interpreter() - - class BlockingStderr: - def read(self): - raise AssertionError("stderr.read() would block if the pipe is still open") - - async def no_completed_responses(_pending_tasks): - return None - - async def stdout_eof(_timeout): - return "" - - interp.deno_process = SimpleNamespace( - stderr=BlockingStderr(), - kill=lambda: None, - poll=lambda: 0, - ) - interp._pending_file_ops = {} - interp._send_completed_responses = no_completed_responses - interp._read_with_timeout_async = stdout_eof - - try: - asyncio.run(interp._execute_async(execute_request_id=1)) - except SandboxFatalError as exc: - assert "Deno subprocess stopped producing stdout" in str(exc) - else: - raise AssertionError("expected SandboxFatalError on Deno stdout EOF") diff --git a/tests/test_iteration_execution_timeout.py b/tests/test_iteration_execution_timeout.py index 4c62ce75..962abe8a 100644 --- a/tests/test_iteration_execution_timeout.py +++ b/tests/test_iteration_execution_timeout.py @@ -1,3 +1,5 @@ +"""Selected iteration deadlines preserve state; failed recovery is bounded.""" + from __future__ import annotations import asyncio @@ -8,19 +10,6 @@ import pytest -class _FakeRepl: - def __init__(self): - self.calls = [] - - def execute(self, code, variables=None, timeout=None): - self.calls.append({"code": code, "variables": variables, "timeout": timeout}) - return "[Success] ok" - - async def aexecute(self, code, variables=None, timeout=None): - self.calls.append({"code": code, "variables": variables, "timeout": timeout}) - return "[Success] ok" - - class _SequentialActions: def __init__(self, *actions: SimpleNamespace) -> None: self.actions = list(actions) @@ -43,257 +32,6 @@ def __getitem__(self, key: str) -> str: return getattr(self, key) -def _build_executor(): - from predict_rlm.predict_rlm import PredictRLM - - executor = PredictRLM.__new__(PredictRLM) - executor.signature = dspy.Signature("question -> answer") - executor.max_iterations = 3 - executor.verbose = False - executor._user_tools = {} - executor.generate_action = MagicMock() - executor._partial_pending_entry = None - executor._partial_history = None - executor._partial_pending_start = None - def _process_execution_result(pred, *args): - result = args[-3] - return { - "result": result, - "pred_code": getattr(pred, "code", None), - } - - executor._process_execution_result = _process_execution_result - return executor - - -def test_missing_action_timeout_preserves_existing_execution_call(): - executor = _build_executor() - executor.generate_action.acall = AsyncMock( - return_value=SimpleNamespace(reasoning="run it", code="print('ok')") - ) - repl = _FakeRepl() - - async def _run(): - return await executor._aexecute_iteration( - repl, - variables=[], - history=MagicMock(), - iteration=0, - input_args={"question": "q"}, - output_field_names=["answer"], - ) - - result = asyncio.run(_run()) - - assert result["result"] == "[Success] ok" - assert repl.calls == [ - { - "code": "print('ok')", - "variables": {"question": "q"}, - "timeout": None, - } - ] - - -def test_positive_action_timeout_is_passed_to_execution(): - executor = _build_executor() - executor.generate_action.acall = AsyncMock( - return_value=SimpleNamespace( - reasoning="run with a cap", - code="print('ok')", - execution_timeout_seconds=2.5, - ) - ) - repl = _FakeRepl() - - async def _run(): - return await executor._aexecute_iteration( - repl, - variables=[], - history=MagicMock(), - iteration=0, - input_args={}, - output_field_names=["answer"], - ) - - asyncio.run(_run()) - - assert repl.calls[0]["timeout"] == 2.5 - - -def test_positive_action_timeout_is_passed_to_sync_execution(): - executor = _build_executor() - executor.generate_action = MagicMock( - return_value=SimpleNamespace( - reasoning="run with a cap", - code="print('ok')", - execution_timeout_seconds=3, - ) - ) - repl = _FakeRepl() - - executor._execute_iteration( - repl, - variables=[], - history=MagicMock(), - iteration=0, - input_args={}, - output_field_names=["answer"], - ) - - assert repl.calls[0]["timeout"] == 3.0 - - -@pytest.mark.parametrize("timeout_value", [None, 1, 2.5]) -def test_action_timeout_accepts_null_and_finite_positive_values(timeout_value): - executor = _build_executor() - pred_kwargs = {"reasoning": "run it", "code": "print('ok')"} - if timeout_value is not None: - pred_kwargs["execution_timeout_seconds"] = timeout_value - pred = SimpleNamespace(**pred_kwargs) - - result = executor._action_execution_timeout(pred) - - assert result == (None if timeout_value is None else float(timeout_value)) - - -def test_action_timeout_uses_shared_validation_helper(monkeypatch): - from predict_rlm import predict_rlm - - calls = [] - - def fake_validate_execution_timeout(value): - calls.append(value) - return 4.0 - - monkeypatch.setattr( - predict_rlm, - "validate_execution_timeout", - fake_validate_execution_timeout, - ) - executor = _build_executor() - pred = SimpleNamespace( - reasoning="run it", - code="print('ok')", - execution_timeout_seconds=4, - ) - - assert executor._action_execution_timeout(pred) == 4.0 - assert calls == [4] - - -@pytest.mark.parametrize( - "timeout_value", - [True, False, "2", 0, -1, float("nan"), float("inf"), -float("inf")], -) -def test_invalid_action_timeout_fails_before_execution(timeout_value): - executor = _build_executor() - executor.generate_action.acall = AsyncMock( - return_value=SimpleNamespace( - reasoning="bad cap", - code="print('ok')", - execution_timeout_seconds=timeout_value, - ) - ) - repl = _FakeRepl() - - async def _run(): - return await executor._aexecute_iteration( - repl, - variables=[], - history=MagicMock(), - iteration=0, - input_args={}, - output_field_names=["answer"], - ) - - with pytest.raises(RuntimeError, match="invalid execution_timeout_seconds"): - asyncio.run(_run()) - assert repl.calls == [] - - -@pytest.mark.parametrize( - ("timeout_value", "expected_log"), - [(None, "Execution timeout: null"), (2.5, "Execution timeout: 2.5s")], -) -def test_verbose_iteration_log_includes_execution_timeout(timeout_value, expected_log, capsys): - executor = _build_executor() - executor.verbose = True - pred_kwargs = {"reasoning": "run it", "code": "print('ok')"} - if timeout_value is not None: - pred_kwargs["execution_timeout_seconds"] = timeout_value - executor.generate_action.acall = AsyncMock(return_value=SimpleNamespace(**pred_kwargs)) - repl = _FakeRepl() - - async def _run(): - return await executor._aexecute_iteration( - repl, - variables=[], - history=MagicMock(), - iteration=1, - input_args={}, - output_field_names=["answer"], - ) - - asyncio.run(_run()) - - captured = capsys.readouterr() - assert expected_log in captured.err - - -@pytest.mark.asyncio -async def test_jspi_per_iteration_timeout_has_recoverable_host_grace(): - from predict_rlm.backends import JspiBackend - from predict_rlm.execution_timeout import ( - DEFAULT_RECOVERABLE_EXECUTION_TIMEOUT_GRACE_SECONDS, - ITERATION_TIMEOUT_FAILURE_CLASS, - format_recoverable_timeout_result, - recoverable_timeout_host_deadline_seconds, - ) - - interpreter = JspiBackend.__new__(JspiBackend) - spans = [] - killed = [] - interpreter._write_telemetry_span = lambda name, **kwargs: spans.append( - {"name": name, **kwargs} - ) - interpreter._telemetry_pending_tool_count = lambda: 0 - interpreter._telemetry_pending_file_ops_count = lambda: 0 - - async def _kill_sandbox(): - killed.append(True) - - interpreter._akill_sandbox = _kill_sandbox - - async def _slow_execute(_request_id): - await asyncio.sleep(1.05) - return format_recoverable_timeout_result( - {"timeout": {"seconds": 0.01}, "stdout": "late\n", "stderr": ""} - ) - - interpreter._execute_async = _slow_execute - - start = asyncio.get_running_loop().time() - result = await interpreter._execute_with_timeout( - 7, - timeout_seconds=0.01, - timeout_failure_class=ITERATION_TIMEOUT_FAILURE_CLASS, - ) - - elapsed = asyncio.get_running_loop().time() - start - - assert DEFAULT_RECOVERABLE_EXECUTION_TIMEOUT_GRACE_SECONDS == 30.0 - assert recoverable_timeout_host_deadline_seconds( - 0.01, - ITERATION_TIMEOUT_FAILURE_CLASS, - ) == 30.01 - assert elapsed >= 1.0 - assert "[Timeout] Iteration execution timed out after 0.01s" in result - assert "[stdout]\nlate" in result - assert killed == [] - assert any(span["name"] == "sandbox.execute.timeout" for span in spans) - - @pytest.mark.asyncio async def test_jspi_silent_iteration_timeout_recovery_failure_is_bounded(monkeypatch): import predict_rlm.execution_timeout as execution_timeout @@ -343,57 +81,6 @@ async def _silent_execute(_request_id): ) -@pytest.mark.asyncio -async def test_jspi_timeout_result_formats_buffered_stdout_and_stderr(): - from predict_rlm.backends import JspiBackend - - interpreter = JspiBackend.__new__(JspiBackend) - interpreter._pending_file_ops = {} - interpreter._active_tool_count = 0 - interpreter._sync_files = lambda: None - interpreter._wait_and_send_all_responses = AsyncMock() - interpreter._send_completed_responses = AsyncMock() - interpreter._read_with_timeout_async = AsyncMock( - return_value='{"jsonrpc":"2.0","result":{"timeout":{"seconds":2.5},' - '"stdout":"out before\\n","stderr":"err before\\n"},"id":9}' - ) - - result = await interpreter._execute_async(9) - - assert "[Timeout] Iteration execution timed out after 2.5s" in result - assert "[stdout]\nout before" in result - assert "[stderr]\nerr before" in result - - -@pytest.mark.integration -def test_jspi_recoverable_timeout_preserves_output_and_globals(): - from predict_rlm.backends import JspiBackend - - interpreter = JspiBackend(preinstall_packages=False, exec_timeout=5.0) - try: - result = interpreter.execute( - """ -import sys -print("stdout before timeout") -print("stderr before timeout", file=sys.stderr) -survived_value = 123 -while True: - pass -""", - timeout=0.2, - ) - - output = str(result) - assert "[Timeout] Iteration execution timed out after 0.2s" in output - assert "stdout before timeout" in output - assert "stderr before timeout" in output - - followup = interpreter.execute("survived_value") - assert followup == 123 - finally: - interpreter.shutdown() - - @pytest.mark.integration def test_predict_rlm_jspi_timeout_preserves_state_history_and_predict_tool(): from predict_rlm import PredictRLM @@ -443,15 +130,11 @@ def test_predict_rlm_jspi_timeout_preserves_state_history_and_predict_tool(): prediction = rlm(prompt="exercise deno timeout recovery") assert prediction.answer == "pre-timeout prediction -> post-timeout prediction / 123" - assert [call["iteration"] for call in actions.calls] == ["1/2", "2/2"] - assert mock_predictor.acall.await_count == 2 - assert [call.kwargs["question"] for call in mock_predictor.acall.await_args_list] == [ - "first call", - "second call", - ] assert len(prediction.trace.steps) == 2 timeout_step, final_step = prediction.trace.steps - assert "[Timeout] Iteration execution timed out after 0.2s" in timeout_step.untruncated_output + assert ( + "[Timeout] Iteration execution timed out after 0.2s" in timeout_step.untruncated_output + ) assert "first predict: pre-timeout prediction" in timeout_step.untruncated_output assert "marker before timeout: 123" in timeout_step.untruncated_output assert final_step.output == ( @@ -462,23 +145,22 @@ def test_predict_rlm_jspi_timeout_preserves_state_history_and_predict_tool(): assert "first predict: pre-timeout prediction" in second_history -@pytest.mark.integration -def test_jspi_no_timeout_execution_still_returns_output_and_stderr(): - from predict_rlm.backends import JspiBackend +@pytest.mark.asyncio +async def test_nonfinite_model_deadline_fails_before_execution(): + from dspy.primitives.repl_types import REPLHistory - interpreter = JspiBackend(preinstall_packages=False, exec_timeout=5.0) - try: - result = interpreter.execute( - """ -import sys -print("stdout ok") -print("stderr ok", file=sys.stderr) -""" - ) - finally: - interpreter.shutdown() + from predict_rlm import PredictRLM - output = str(result) - assert "stdout ok" in output - assert "stderr ok" in output - assert "[Timeout]" not in output + rlm = PredictRLM("question -> answer", model_execution_timeout=True) + rlm.generate_action = MagicMock() + rlm.generate_action.acall = AsyncMock( + return_value=dspy.Prediction( + reasoning="invalid cap", + code="while True: pass", + execution_timeout_seconds=float("nan"), + ) + ) + repl = MagicMock() + repl.aexecute = AsyncMock(side_effect=AssertionError("unbounded code executed")) + with pytest.raises(RuntimeError, match="invalid execution_timeout_seconds"): + await rlm._aexecute_iteration(repl, [], REPLHistory(), 0, {}, ["answer"]) diff --git a/tests/test_iteration_usage.py b/tests/test_iteration_usage.py deleted file mode 100644 index e8a70499..00000000 --- a/tests/test_iteration_usage.py +++ /dev/null @@ -1,62 +0,0 @@ -"""Per-iteration token usage/cost accounting on IterationStep.usage.""" - -from __future__ import annotations - -import dspy - -from predict_rlm import PredictRLM -from predict_rlm.trace import IterationStep, LMUsage, PredictCallGroup, TokenUsage - - -def test_iteration_step_usage_defaults_to_empty_lm_usage(): - step = IterationStep( - iteration=1, - reasoning="r", - code="c", - output="o", - untruncated_output="o", - duration_ms=1, - ) - assert isinstance(step.usage, LMUsage) - assert step.usage.main.input_tokens == 0 - assert step.usage.sub.cost == 0.0 - - -def test_build_iteration_usage_combines_main_and_sub(): - rlm = PredictRLM("q -> a", sub_lm=dspy.LM("openai/gpt-4o", api_key="x")) - # Main action-LM usage stashed by _record_action_generation_ok. - rlm._last_action_lm_usage = TokenUsage(input_tokens=2000, output_tokens=100, cost=0.012) - predict_calls = [ - PredictCallGroup( - signature="x -> y", - instructions=None, - model="openai/gpt-4o", - total_usage=TokenUsage(input_tokens=50, output_tokens=10, cost=0.001), - calls=[], - ), - PredictCallGroup( - signature="x -> z", - instructions=None, - model="openai/gpt-4o", - total_usage=TokenUsage(input_tokens=30, output_tokens=5, cost=0.0005), - calls=[], - ), - ] - - usage = rlm._build_iteration_usage(predict_calls) - - assert usage.main.input_tokens == 2000 - assert usage.main.output_tokens == 100 - assert usage.main.cost == 0.012 - assert usage.sub.input_tokens == 80 # 50 + 30 - assert usage.sub.output_tokens == 15 - assert round(usage.sub.cost, 4) == 0.0015 - # Stash is consumed so it cannot leak into the next iteration. - assert rlm._last_action_lm_usage is None - - -def test_build_iteration_usage_without_stash_is_empty_main(): - rlm = PredictRLM("q -> a", sub_lm=dspy.LM("openai/gpt-4o", api_key="x")) - usage = rlm._build_iteration_usage([]) - assert usage.main.input_tokens == 0 - assert usage.sub.input_tokens == 0 diff --git a/tests/test_jspi_async_operations.py b/tests/test_jspi_async_operations.py index 6c56852e..83e5d822 100644 --- a/tests/test_jspi_async_operations.py +++ b/tests/test_jspi_async_operations.py @@ -1,10 +1,8 @@ from __future__ import annotations import asyncio -import inspect -import json from types import SimpleNamespace -from unittest.mock import AsyncMock, MagicMock +from unittest.mock import AsyncMock import pytest @@ -13,178 +11,14 @@ from predict_rlm.backends.jspi.execution import JspiExecutionBackend from predict_rlm.execution_timeout import ITERATION_TIMEOUT_FAILURE_CLASS from predict_rlm.runtime import ExecutionSpec -from predict_rlm.workspace import WorkspaceFileInfo - - -def _fail_sync_helper(*args, **kwargs): - raise AssertionError("async JSPI path called a synchronous helper") @pytest.fixture def interpreter() -> JspiBackend: backend = JspiBackend(preinstall_packages=False) - backend._ensure_deno_process = _fail_sync_helper # type: ignore[method-assign] - backend._send_request = _fail_sync_helper # type: ignore[method-assign] - backend._mount_files = _fail_sync_helper # type: ignore[method-assign] - backend._register_tools = _fail_sync_helper # type: ignore[method-assign] return backend -@pytest.mark.asyncio -async def test_aexecute_uses_only_async_setup_helpers(interpreter: JspiBackend): - interpreter._aensure_deno_process = AsyncMock() # type: ignore[method-assign] - interpreter._amount_files = AsyncMock() # type: ignore[method-assign] - interpreter._aregister_tools = AsyncMock() # type: ignore[method-assign] - interpreter._write_stdin_async = AsyncMock() # type: ignore[method-assign] - interpreter._execute_with_timeout = AsyncMock(return_value="ok") # type: ignore[method-assign] - - result = await interpreter._aexecute_inner("print('ok')", {}) - - assert result == "ok" - interpreter._aensure_deno_process.assert_awaited_once_with() - interpreter._amount_files.assert_awaited_once_with() - interpreter._aregister_tools.assert_awaited_once_with() - - -@pytest.mark.asyncio -async def test_aexecute_broken_pipe_uses_async_kill(interpreter: JspiBackend): - interpreter._aensure_deno_process = AsyncMock() # type: ignore[method-assign] - interpreter._amount_files = AsyncMock() # type: ignore[method-assign] - interpreter._aregister_tools = AsyncMock() # type: ignore[method-assign] - interpreter._write_stdin_async = AsyncMock( # type: ignore[method-assign] - side_effect=BrokenPipeError - ) - interpreter._kill_sandbox = _fail_sync_helper # type: ignore[method-assign] - interpreter._akill_sandbox = AsyncMock() # type: ignore[method-assign] - - with pytest.raises(SandboxFatalError, match="BrokenPipeError"): - await interpreter._aexecute_inner("print('ok')", {}) - - interpreter._akill_sandbox.assert_awaited_once_with() - - -@pytest.mark.asyncio -async def test_async_control_operations_use_async_rpc(interpreter: JspiBackend): - interpreter._aensure_deno_process = AsyncMock() # type: ignore[method-assign] - interpreter._asend_request = AsyncMock( # type: ignore[method-assign] - side_effect=[ - {"result": {"mounted": "/sandbox/input.txt"}}, - {"result": {"created": "/sandbox/output"}}, - {"result": {"files": ["/sandbox/output/result.txt"]}}, - { - "result": { - "files": { - "result.txt": { - "type": "file", - "sha256": "abc", - "size": 3, - } - } - } - }, - {"result": {"ok": True}}, - ] - ) - - await interpreter.amount_file_at("/host/input.txt", "/sandbox/input.txt") - await interpreter.amkdir_p("/sandbox/output") - files = await interpreter.alist_dir("/sandbox/output") - manifest = await interpreter.aworkspace_manifest("/sandbox/output") - await interpreter.async_file_to("/sandbox/output/result.txt", "/host/result.txt") - - assert files == ["/sandbox/output/result.txt"] - assert manifest == { - "result.txt": WorkspaceFileInfo(type="file", sha256="abc", size=3) - } - assert [call.args[0] for call in interpreter._asend_request.await_args_list] == [ - "mount_file", - "mkdir_p", - "list_dir", - "workspace_manifest", - "sync_file", - ] - - -@pytest.mark.asyncio -async def test_async_package_setup_uses_async_rpc(interpreter: JspiBackend): - interpreter._aensure_deno_process = AsyncMock() # type: ignore[method-assign] - interpreter._asend_request = AsyncMock( # type: ignore[method-assign] - return_value={"result": {"installed": ["openpyxl"]}} - ) - - await interpreter.aensure_skill_packages(["openpyxl", "openpyxl"]) - - interpreter._asend_request.assert_awaited_once_with( - "install_packages", - {"packages": ["openpyxl"]}, - "installing skill packages", - ) - - -@pytest.mark.asyncio -async def test_async_ready_uses_native_async_startup(interpreter: JspiBackend): - interpreter._aensure_deno_process = AsyncMock() # type: ignore[method-assign] - - await interpreter.aensure_ready() - await interpreter.astart() - - assert interpreter._aensure_deno_process.await_count == 2 - - -def test_async_lifecycle_methods_do_not_use_to_thread(): - methods = ( - JspiBackend._aensure_deno_process, - JspiBackend._ahealth_check, - JspiBackend._asend_request, - JspiBackend._amount_files, - JspiBackend._aregister_tools, - JspiBackend._async_files, - JspiBackend._akill_sandbox, - JspiBackend.aensure_ready, - JspiBackend.astart, - JspiBackend.aensure_skill_packages, - JspiBackend.amount_file_at, - JspiBackend.amkdir_p, - JspiBackend.alist_dir, - JspiBackend.aworkspace_manifest, - JspiBackend.async_file_to, - JspiBackend.aexecute, - JspiBackend.ainterrupt, - JspiBackend.ashutdown, - ) - - for method in methods: - source = inspect.getsource(method) - assert "to_thread" not in source, method.__qualname__ - - -@pytest.mark.asyncio -async def test_ainterrupt_does_not_call_sync_interrupt(interpreter: JspiBackend): - interpreter.interrupt = _fail_sync_helper # type: ignore[method-assign] - interpreter._akill_sandbox = AsyncMock() # type: ignore[method-assign] - - await interpreter.ainterrupt() - - interpreter._akill_sandbox.assert_awaited_once_with() - - -@pytest.mark.asyncio -async def test_cancel_execution_quiesces_without_killing_sandbox(interpreter: JspiBackend): - process = SimpleNamespace(poll=lambda: None, send_signal=MagicMock()) - interpreter.deno_process = process - interpreter._active_execute_request_id = 7 - interpreter._execute_async = AsyncMock( # type: ignore[method-assign] - side_effect=SandboxExecutionError("KeyboardInterrupt") - ) - interpreter._akill_sandbox = AsyncMock() # type: ignore[method-assign] - - await interpreter.acancel_execution() - - process.send_signal.assert_called_once() - interpreter._execute_async.assert_awaited_once_with(7) - interpreter._akill_sandbox.assert_not_awaited() - - @pytest.mark.asyncio async def test_cancel_execution_retries_interrupt_until_execution_quiesces( interpreter: JspiBackend, @@ -218,40 +52,6 @@ async def finish_execution(request_id): interpreter._akill_sandbox.assert_not_awaited() -@pytest.mark.asyncio -async def test_ashutdown_does_not_call_sync_shutdown(interpreter: JspiBackend): - interpreter.shutdown = _fail_sync_helper # type: ignore[method-assign] - interpreter.deno_process = type( - "Process", - (), - {"poll": lambda self: None, "stdin": None}, - )() - interpreter._write_stdin_async = AsyncMock() # type: ignore[method-assign] - interpreter._await_process_exit = AsyncMock(return_value=True) # type: ignore[method-assign] - - await asyncio.wait_for(interpreter.ashutdown(), timeout=0.1) - - message = json.loads(interpreter._write_stdin_async.await_args.args[0]) - assert message["method"] == "shutdown" - assert interpreter.deno_process is None - - -@pytest.mark.asyncio -async def test_async_rpc_uses_only_async_fd_primitives(interpreter: JspiBackend): - interpreter._write_stdin_async = AsyncMock() # type: ignore[method-assign] - interpreter._read_with_timeout_async = AsyncMock( # type: ignore[method-assign] - side_effect=lambda timeout: json.dumps( - {"jsonrpc": "2.0", "id": interpreter._request_id, "result": {"ok": True}} - ) - ) - - response = await interpreter._asend_request("mkdir_p", {"path": "/sandbox/x"}, "test") - - assert response["result"] == {"ok": True} - interpreter._write_stdin_async.assert_awaited_once() - interpreter._read_with_timeout_async.assert_awaited_once() - - @pytest.mark.asyncio async def test_aexecute_skips_post_hooks_after_fatal_failure(interpreter: JspiBackend): hook = AsyncMock() @@ -287,30 +87,6 @@ async def block(code, variables): hook.assert_not_awaited() -@pytest.mark.asyncio -async def test_sync_worker_quarantine_defers_jspi_shutdown(interpreter: JspiBackend): - release = __import__("threading").Event() - - worker = interpreter._start_sync_worker(lambda: release.wait()) - shutdown = MagicMock() - interpreter.shutdown = shutdown # type: ignore[method-assign] - - retired = interpreter.retire_when_sync_workers_finish() - await asyncio.sleep(0) - - assert retired is True - shutdown.assert_not_called() - - release.set() - await asyncio.wait_for(worker.wait(), timeout=1) - for _ in range(100): - if shutdown.called: - break - await asyncio.sleep(0.01) - - shutdown.assert_called_once_with() - - @pytest.mark.asyncio async def test_aexecute_preserves_primary_error_when_post_hook_fails( interpreter: JspiBackend, diff --git a/tests/test_jspi_workspace_lifecycle.py b/tests/test_jspi_workspace_lifecycle.py deleted file mode 100644 index a36774af..00000000 --- a/tests/test_jspi_workspace_lifecycle.py +++ /dev/null @@ -1,77 +0,0 @@ -from __future__ import annotations - -import asyncio -import shutil -from pathlib import Path -from unittest.mock import AsyncMock, MagicMock - -import dspy -import pytest - -from predict_rlm import PredictRLM, Workspace - - -class WorkspaceCancellationSignature(dspy.Signature): - workspace: Workspace = dspy.InputField() - answer: str = dspy.OutputField() - - -@pytest.mark.integration -@pytest.mark.skipif(shutil.which("deno") is None, reason="JSPI lifecycle test requires Deno") -@pytest.mark.asyncio -async def test_jspi_cancellation_flushes_mirror_before_sandbox_shutdown(tmp_path: Path): - mutation_completed = asyncio.Event() - - async def signal_mutation() -> str: - """Tell the host that the sandbox mutation completed.""" - asyncio.get_running_loop().call_later(0.1, mutation_completed.set) - return "ok" - - workspace_root = tmp_path / "workspace" - workspace_root.mkdir() - source = workspace_root / "source.txt" - source.write_text("before", encoding="utf-8") - - rlm = PredictRLM( - WorkspaceCancellationSignature, - lm=MagicMock(history=[]), - tools={"signal_mutation": signal_mutation}, - max_iterations=1, - verbose=False, - ) - rlm.generate_action.acall = AsyncMock( - return_value=dspy.Prediction( - reasoning="mutate the workspace and remain active", - code=( - "from pathlib import Path\n" - "Path('/sandbox/workspace/source.txt').write_text('after-cancel')\n" - "await signal_mutation()\n" - "while True:\n" - " pass" - ), - ) - ) - rlm._configure_run_predictors = MagicMock() - - invocation = asyncio.create_task( - rlm.aforward(workspace=Workspace(path=str(workspace_root))) - ) - mutation_wait = asyncio.create_task(mutation_completed.wait()) - done, _ = await asyncio.wait( - {invocation, mutation_wait}, - timeout=30, - return_when=asyncio.FIRST_COMPLETED, - ) - if invocation in done: - await invocation - if mutation_wait not in done: - invocation.cancel() - mutation_wait.cancel() - await asyncio.gather(invocation, mutation_wait, return_exceptions=True) - pytest.fail("sandbox mutation did not complete before the test timeout") - invocation.cancel() - - with pytest.raises(asyncio.CancelledError): - await asyncio.wait_for(invocation, timeout=30) - - assert source.read_text(encoding="utf-8") == "after-cancel" diff --git a/tests/test_lm_config.py b/tests/test_lm_config.py index 574bee44..8cf210f9 100644 --- a/tests/test_lm_config.py +++ b/tests/test_lm_config.py @@ -13,45 +13,6 @@ def _skip_env_validation(monkeypatch): monkeypatch.setattr(lm_config, "validate_lm_env", lambda _lm: None) -def test_kimi_reasoning_effort_none_disables_native_thinking(monkeypatch): - _skip_env_validation(monkeypatch) - - config = lm_config.get_lm_config("moonshot/kimi-k2.6", reasoning_effort="none") - - assert config["extra_body"] == {"thinking": {"type": "disabled"}} - assert "reasoning_effort" not in config - - -def test_kimi_reasoning_effort_enables_native_thinking_without_fake_effort(monkeypatch): - _skip_env_validation(monkeypatch) - - config = lm_config.get_lm_config("moonshot/kimi-k2.6", reasoning_effort="low") - - assert config["extra_body"] == {"thinking": {"type": "enabled"}} - assert "reasoning_effort" not in config - - -def test_kimi_unspecified_reasoning_leaves_provider_default(monkeypatch): - _skip_env_validation(monkeypatch) - - config = lm_config.get_lm_config("moonshot/kimi-k2.6", reasoning_effort=None) - - assert "extra_body" not in config - assert "reasoning_effort" not in config - - -def test_non_kimi_reasoning_effort_behavior_is_unchanged(monkeypatch): - _skip_env_validation(monkeypatch) - - config = lm_config.get_lm_config("openai/gpt-5.4", reasoning_effort="low") - disabled_config = lm_config.get_lm_config("openai/gpt-5.4", reasoning_effort="none") - - assert config["reasoning_effort"] == "low" - assert "reasoning_effort" not in disabled_config - assert "extra_body" not in config - assert "extra_body" not in disabled_config - - def test_build_lm_retries_litellm_rate_limits_with_tenacity(monkeypatch): _skip_env_validation(monkeypatch) monkeypatch.setattr(lm_config, "_RATE_LIMIT_RETRY_WAIT", wait_none(), raising=False) diff --git a/tests/test_logging.py b/tests/test_logging.py deleted file mode 100644 index 3cbc9b9e..00000000 --- a/tests/test_logging.py +++ /dev/null @@ -1,158 +0,0 @@ -import logging - -import predict_rlm._logging as logging_module -from predict_rlm._logging import ( - DEBUG_HANDLER_MARKER, - PACKAGE_LOGGER_NAME, - TRACE_HANDLER_MARKER, - TRACE_LOGGER_NAME, - _PredictRLMDebugFormatter, - configure_predict_rlm_logging, -) - - -def _snapshot_logger(logger: logging.Logger): - return ( - logger.level, - logger.propagate, - logger.disabled, - list(logger.handlers), - ) - - -def _restore_logger(logger: logging.Logger, state) -> None: - level, propagate, disabled, handlers = state - logger.setLevel(level) - logger.propagate = propagate - logger.disabled = disabled - logger.handlers[:] = handlers - - -def _format_debug_message(message: str) -> str: - record = logging.LogRecord( - "predict_rlm.test", - logging.DEBUG, - __file__, - 1, - message, - (), - None, - ) - formatter = _PredictRLMDebugFormatter("%(levelname)s:%(message)s") - return formatter.format(record) - - -def test_debug_formatter_colors_error_events_red(): - formatted = _format_debug_message("rlm.execute.error error_type=ValueError") - - assert formatted.startswith("\033[31m") - assert formatted.endswith("\033[0m") - - -def test_debug_formatter_colors_status_error_records_red(): - formatted = _format_debug_message("sbx.runner.exited status=error") - - assert formatted.startswith("\033[31m") - assert formatted.endswith("\033[0m") - - -def test_debug_formatter_leaves_non_error_records_plain(): - formatted = _format_debug_message("rlm.execute.ok duration_ms=12") - - assert formatted == "DEBUG:rlm.execute.ok duration_ms=12" - - -def test_debug_logging_can_restore_prior_logger_state(): - logger = logging.getLogger(PACKAGE_LOGGER_NAME) - original_logger = _snapshot_logger(logger) - original_module_state = logging_module._debug_logger_state - try: - logging_module._debug_logger_state = None - logger.handlers[:] = [] - logger.setLevel(logging.WARNING) - logger.propagate = False - logger.disabled = False - - configure_predict_rlm_logging(debug=True) - - assert logger.level == logging.DEBUG - assert any(getattr(handler, DEBUG_HANDLER_MARKER, False) for handler in logger.handlers) - - configure_predict_rlm_logging(debug=False) - - assert logger.level == logging.WARNING - assert logger.propagate is False - assert logger.disabled is False - assert not any( - getattr(handler, DEBUG_HANDLER_MARKER, False) for handler in logger.handlers - ) - finally: - logging_module._debug_logger_state = original_module_state - _restore_logger(logger, original_logger) - - -def test_verbose_trace_detail_does_not_hard_wrap_long_logical_lines(): - long_line = "x" * 160 - - rendered = logging_module._render_trace_detail("output:", long_line) - - assert long_line in rendered - - -def test_verbose_trace_header_colors_only_current_iteration_green(): - rendered = logging_module._render_trace_header(3, 12) - - assert rendered == ( - "\033[1;97mRLM turn " - "\033[1;32m3" - "\033[1;97m/12" - "\033[0m" - ) - - -def test_verbose_trace_detail_uses_lighter_label_and_dim_body(): - rendered = logging_module._render_trace_detail("output:", "result") - - assert "\033[3;97moutput:\033[0m" in rendered - assert "\033[2mresult\033[0m" in rendered - - -def test_verbose_trace_detail_preserves_code_syntax_color(): - rendered = logging_module._render_trace_detail( - "code:", - "value = {'x': 1}", - syntax="python", - ) - - assert "\033[33m'" in rendered - assert "\033[34m1" in rendered - - -def test_verbose_logging_can_restore_prior_trace_logger_state(): - logger = logging.getLogger(TRACE_LOGGER_NAME) - original_logger = _snapshot_logger(logger) - original_module_state = logging_module._trace_logger_state - try: - logging_module._trace_logger_state = None - logger.handlers[:] = [] - logger.setLevel(logging.WARNING) - logger.propagate = True - logger.disabled = False - - configure_predict_rlm_logging(verbose=True) - - assert logger.level == logging.INFO - assert logger.propagate is False - assert any(getattr(handler, TRACE_HANDLER_MARKER, False) for handler in logger.handlers) - - configure_predict_rlm_logging(verbose=False) - - assert logger.level == logging.WARNING - assert logger.propagate is True - assert logger.disabled is False - assert not any( - getattr(handler, TRACE_HANDLER_MARKER, False) for handler in logger.handlers - ) - finally: - logging_module._trace_logger_state = original_module_state - _restore_logger(logger, original_logger) diff --git a/tests/test_predict_rlm.py b/tests/test_predict_rlm.py index aed0039d..75c8d754 100644 --- a/tests/test_predict_rlm.py +++ b/tests/test_predict_rlm.py @@ -1,54 +1,19 @@ -"""Tests for PredictRLM with predict tool for DSPy signatures.""" +"""Core predict output, action-loop, and submit confirmation contracts.""" import asyncio -import hashlib -import logging -import re +import json from typing import Any from unittest.mock import AsyncMock, MagicMock, patch import dspy import pytest from dspy.primitives.code_interpreter import FinalOutput -from dspy.primitives.repl_types import REPLEntry, REPLHistory +from dspy.utils.dummies import DummyLM from pydantic import BaseModel, Field, ValidationError from predict_rlm import PredictRLM from predict_rlm.predict_rlm import _models_from_schema -from predict_rlm.rlm_skills import Skill -from predict_rlm.telemetry import TelemetryContext, classify_failure -from predict_rlm.trace import ( - LMFinishMetadata, - drain_predict_calls, - init_predict_call_collector, - lm_completion_metadata_since, -) - - -def _run(coro): - """Run async predict call from sync test.""" - import nest_asyncio - - nest_asyncio.apply() - loop = asyncio.get_event_loop() - return loop.run_until_complete(coro) - - -def _log_messages(caplog, logger_name: str) -> str: - return "\n".join( - record.getMessage() - for record in caplog.records - if record.name.startswith(logger_name) - ) - - -def _assert_raw_verbose_output(output: str) -> None: - assert "[INFO]" not in output - assert "predict_rlm.trace" not in output - - -def _strip_ansi(output: str) -> str: - return re.sub(r"\x1b\[[0-9;]*m", "", output) +from predict_rlm.trace import PredictCallGroup, TokenUsage class ImageAnalysisSignature(dspy.Signature): @@ -59,38 +24,6 @@ class ImageAnalysisSignature(dspy.Signature): answer: str = dspy.OutputField(desc="Answer to the query") -class DefaultAnswerSignature(dspy.Signature): - """Return an optional answer.""" - - instruction: str = dspy.InputField() - answer: str | None = dspy.OutputField(default=None) - - -class MockLM: - """Mock LM that returns predictable responses for testing.""" - - def __init__(self, responses: dict[str, str] | None = None): - self.responses = responses or {} - self.calls: list[dict] = [] - - def __call__(self, messages=None, **kwargs): - self.calls.append({"messages": messages, "kwargs": kwargs}) - content = messages[-1].get("content", "") if messages else "" - prompt = content if isinstance(content, str) else str(content) - for key, response in self.responses.items(): - if key.lower() in prompt.lower(): - return [response] - return ["Default LM response"] - - -class ListTelemetrySink: - def __init__(self): - self.records: list[dict[str, Any]] = [] - - def write(self, record: dict[str, Any]) -> None: - self.records.append(record) - - class FakeSubmitRepl: def __init__(self, final_payload: dict[str, Any] | None = None): self.final_payload = final_payload or {"answer": "done"} @@ -127,295 +60,275 @@ def __exit__(self, *_args): return False -@pytest.mark.sbx -class TestBackendNameSelection: - """Tests for PredictRLM sandbox backend selection. +@pytest.mark.integration +def test_predict_reconstructs_nested_sandbox_models_and_serializes_items_field(): + from predict_rlm.backends import JspiBackend - Marked ``sbx`` because several cases import SBX symbols (SbxConfig/SbxPool), - which resolve the lazy SBX backend import and require the [sbx] extra. - """ + payload = {"items": [{"name": "repair", "priority": "high", "note": None}]} + interpreter = JspiBackend(preinstall_packages=False) + rlm = PredictRLM( + "query -> answer", + interpreter=interpreter, + sub_lm=DummyLM([{"items": [payload]}]), + max_iterations=1, + ) + code = """ +import json +from typing import Literal +from pydantic import BaseModel + +class Item(BaseModel): + name: str + priority: Literal["high", "low"] + note: str | None = None + +class Order(BaseModel): + items: list[Item] + +extracted = await predict("text: str -> items: list[Order]", text="repair urgently") +SUBMIT(answer=json.dumps({"items": [order.model_dump() for order in extracted.items]})) +""" + try: + with dspy.context(lm=DummyLM([{"reasoning": "extract", "code": code}])): + result = rlm(query="extract repairs") + assert json.loads(result.answer) == {"items": [payload]} + finally: + interpreter.shutdown() + + +def test_reconstructed_defaulted_collections_remain_non_nullable(): + class Item(BaseModel): + name: str + tags: list[str] = Field(default_factory=list) + note: str | None = None + + Model = _models_from_schema(Item.model_json_schema())["Item"] + assert Model(name="x").model_dump() == {"name": "x", "tags": [], "note": None} + with pytest.raises(ValidationError): + Model(name="x", tags=None) + + +@pytest.mark.asyncio +async def test_predict_rejects_non_nullable_null_from_custom_predictor(): + rlm = PredictRLM("question -> answer", sub_lm=DummyLM([])) + # A custom predictor can bypass adapter validation; the tool must still reject null. + with patch("predict_rlm.predict_rlm.dspy.Predict") as predictor: + predictor.return_value.acall = AsyncMock(return_value=dspy.Prediction(items=None)) + with pytest.raises(RuntimeError, match="None for non-Optional"): + await rlm.tools["predict"].func("text: str -> items: list[str]", text="input") + + +@pytest.mark.asyncio +async def test_predict_requires_an_available_lm(): + rlm = PredictRLM("question -> answer") + with dspy.context(lm=None): + with pytest.raises(RuntimeError, match="No LM available"): + await rlm.tools["predict"].func("question -> answer", question="test") + + +@pytest.mark.asyncio +async def test_concurrent_predict_calls_isolate_context_lms_and_restore_caller(): + rlm = PredictRLM("question -> answer") + caller = DummyLM([{"answer": "caller"}]) + + async def run(answer): + with dspy.context(lm=DummyLM([{"answer": answer}])): + await asyncio.sleep(0) + return await rlm.tools["predict"].func("question -> answer", question="same") + + with dspy.context(lm=caller): + results = await asyncio.gather(run("first"), run("second")) + assert results == [{"answer": "first"}, {"answer": "second"}] + assert dspy.settings.lm is caller + assert await rlm.tools["predict"].func("question -> answer", question="same") == { + "answer": "caller" + } - def test_default_backend_remains_jspi(self): - rlm = PredictRLM(ImageAnalysisSignature, sub_lm=MagicMock(), max_iterations=1) - execution_tools = {"predict": MagicMock()} - with patch("predict_rlm.predict_rlm.JspiBackend") as mock_jspi: - mock_repl = MagicMock() - mock_jspi.return_value = mock_repl +@pytest.mark.integration +@pytest.mark.asyncio +async def test_async_loop_executes_actions_recovers_errors_and_submits_output(): + from predict_rlm.backends import JspiBackend - with rlm._interpreter_context(execution_tools=execution_tools) as repl: - assert repl is mock_repl + interpreter = JspiBackend(preinstall_packages=False) + lm = DummyLM( + [ + { + "reasoning": "probe", + "code": "```repl\nvalue = 6\nprint('before failure')\nraise ValueError('retry')\n```", + }, + {"reasoning": "recover", "code": "SUBMIT(answer=str(value * 7))"}, + ] + ) + try: + rlm = PredictRLM("query -> answer", interpreter=interpreter, max_iterations=2) + with dspy.context(lm=lm): + result = await rlm.acall(query="compute") + assert result.answer == "42" + assert result.trace.status == "completed" + assert "before failure" in result.trace.steps[0].untruncated_output + assert result.trace.steps[0].error is True + finally: + interpreter.shutdown() - mock_jspi.assert_called_once() - assert mock_jspi.call_args.kwargs["tools"] == execution_tools - def test_explicit_sbx_backend_uses_sbx_interpreter(self): - from predict_rlm import BackendName, SbxConfig +class TestSubmitConfirmation: + """Tests for configurable submit confirmation in the main RLM loop.""" - rlm = PredictRLM( - ImageAnalysisSignature, - sub_lm=MagicMock(), - max_iterations=1, - sandbox_backend=BackendName.SBX, - sbx_config=SbxConfig(name="test-sbx"), - ) - execution_tools = {"predict": MagicMock()} - - with patch("predict_rlm.backends.sbx.SbxBackend") as mock_sbx: - mock_repl = MagicMock() - mock_sbx.return_value = mock_repl - - with rlm._interpreter_context(execution_tools=execution_tools) as repl: - assert repl is mock_repl - - mock_sbx.assert_called_once() - assert mock_sbx.call_args.kwargs["tools"] == execution_tools - assert mock_sbx.call_args.kwargs["config"].name == "test-sbx" - - def test_custom_interpreter_conflicts_with_explicit_backend(self): - with pytest.raises(ValueError, match="interpreter.*sandbox_backend"): - PredictRLM( - ImageAnalysisSignature, - sub_lm=MagicMock(), - max_iterations=1, - interpreter=MagicMock(), - sandbox_backend="sbx", - ) + @staticmethod + def _prediction(code: str, reasoning: str = "thinking") -> dspy.Prediction: + return dspy.Prediction(reasoning=reasoning, code=code) - def test_debug_configures_injected_interpreter_debug(self): - interpreter = MagicMock() - interpreter.configure_debug = MagicMock() - interpreter.shutdown = MagicMock() - rlm = PredictRLM( - ImageAnalysisSignature, - sub_lm=MagicMock(), - max_iterations=1, - interpreter=interpreter, - debug=True, - ) + def _run_sync( + self, + rlm: PredictRLM, + actions: list[dspy.Prediction], + repl: FakeSubmitRepl | None = None, + ) -> dspy.Prediction: + repl = repl or FakeSubmitRepl() + mock_lm = MagicMock() + mock_lm.history = [] + rlm.generate_action = MagicMock(side_effect=actions) - with rlm._interpreter_context(execution_tools={"predict": MagicMock()}) as repl: - assert repl is interpreter - interpreter.configure_debug.assert_called_once_with(True) + with ( + dspy.context(lm=mock_lm), + patch.object( + rlm, "_interpreter_context", return_value=FakeInterpreterContext(repl) + ), + ): + return rlm._forward_traced(None, images=["img"], query="Original task") - interpreter.shutdown.assert_not_called() + async def _run_async( + self, + rlm: PredictRLM, + actions: list[dspy.Prediction], + repl: FakeSubmitRepl | None = None, + ) -> dspy.Prediction: + repl = repl or FakeSubmitRepl() + mock_lm = MagicMock() + mock_lm.history = [] + rlm.generate_action = MagicMock() + rlm.generate_action.acall = AsyncMock(side_effect=actions) - def test_verbose_configures_injected_interpreter_verbose(self): - interpreter = MagicMock() - interpreter.configure_verbose = MagicMock() - interpreter.shutdown = MagicMock() - rlm = PredictRLM( - ImageAnalysisSignature, - sub_lm=MagicMock(), - max_iterations=1, - interpreter=interpreter, - verbose=True, - ) + with ( + dspy.context(lm=mock_lm), + patch.object( + rlm, "_interpreter_context", return_value=FakeInterpreterContext(repl) + ), + ): + return await rlm._aforward_traced(None, images=["img"], query="Original task") - with rlm._interpreter_context(execution_tools={"predict": MagicMock()}) as repl: - assert repl is interpreter - interpreter.configure_verbose.assert_called_once_with(True) + def test_first_submit_prompts_and_second_submit_completes(self): + seen_contexts = [] - interpreter.shutdown.assert_not_called() + def confirm(context): + seen_contexts.append(context) + return "Please verify the answer before final submit." - def test_default_verbose_configures_injected_interpreter_verbose(self): - interpreter = MagicMock() - interpreter.configure_verbose = MagicMock() - interpreter.shutdown = MagicMock() rlm = PredictRLM( ImageAnalysisSignature, sub_lm=MagicMock(), - max_iterations=1, - interpreter=interpreter, + max_iterations=3, + submit_confirmation=confirm, ) - with rlm._interpreter_context(execution_tools={"predict": MagicMock()}) as repl: - assert repl is interpreter - interpreter.configure_verbose.assert_called_once_with(True) - - interpreter.shutdown.assert_not_called() - - def test_explicit_false_verbose_configures_injected_interpreter_quiet(self): - interpreter = MagicMock() - interpreter.configure_verbose = MagicMock() - interpreter.shutdown = MagicMock() - rlm = PredictRLM( - ImageAnalysisSignature, - sub_lm=MagicMock(), - max_iterations=1, - interpreter=interpreter, - verbose=False, + repl = FakeSubmitRepl() + result = self._run_sync( + rlm, + [ + self._prediction("SUBMIT(answer='done')", reasoning="first submit"), + self._prediction("SUBMIT(answer='done')", reasoning="second submit"), + ], + repl=repl, ) - with rlm._interpreter_context(execution_tools={"predict": MagicMock()}) as repl: - assert repl is interpreter - interpreter.configure_verbose.assert_called_once_with(False) - - interpreter.shutdown.assert_not_called() + assert result.answer == "done" + assert result.trace.status == "completed" + assert [step.output for step in result.trace.steps] == [ + "Please verify the answer before final submit.", + "FINAL: {'answer': 'done'}", + ] + assert len(seen_contexts) == 1 + context = seen_contexts[0] + assert context.inputs == {"images": ["img"], "query": "Original task"} + assert context.submitted_payload == {"answer": "done"} - def test_configures_injected_interpreter_runtime_logging(self): - class Interpreter: - def __init__(self) -> None: - self.tools = {} - self.output_fields = [] - self.runtime_kwargs = None - self.shutdown = MagicMock() + def test_non_submit_after_confirmation_clears_pending_confirmation(self): + prompts = [] - def configure_runtime(self, **kwargs): - self.runtime_kwargs = kwargs + def confirm(context): + prompts.append(context.iteration) + return f"confirm attempt {context.iteration}" - interpreter = Interpreter() rlm = PredictRLM( ImageAnalysisSignature, sub_lm=MagicMock(), - max_iterations=1, - interpreter=interpreter, - debug=True, - verbose=True, + max_iterations=5, + submit_confirmation=confirm, ) - with rlm._interpreter_context(execution_tools={"predict": MagicMock()}) as repl: - assert repl is interpreter - - assert interpreter.runtime_kwargs == {"debug": True, "verbose": True} - interpreter.shutdown.assert_not_called() - - def test_injected_interpreter_logging_does_not_mutate_ad_hoc_attrs(self): - class Interpreter: - def __init__(self) -> None: - self.tools = {} - self.output_fields = [] - self.debug = False - self._debug = False - self.verbose = False - self._verbose = False - self.shutdown = MagicMock() - - interpreter = Interpreter() - rlm = PredictRLM( - ImageAnalysisSignature, - sub_lm=MagicMock(), - max_iterations=1, - interpreter=interpreter, - debug=True, - verbose=True, + result = self._run_sync( + rlm, + [ + self._prediction("SUBMIT(answer='done')"), + self._prediction("print('checking')"), + self._prediction("SUBMIT(answer='done')"), + self._prediction("SUBMIT(answer='done')"), + ], ) - with rlm._interpreter_context(execution_tools={"predict": MagicMock()}) as repl: - assert repl is interpreter - - assert interpreter.debug is False - assert interpreter._debug is False - assert interpreter.verbose is False - assert interpreter._verbose is False - interpreter.shutdown.assert_not_called() - - def test_sbx_pool_requires_sbx_backend(self): - from predict_rlm import SbxPool - - with pytest.raises(ValueError, match="sbx_pool.*sandbox_backend='sbx'"): - PredictRLM( - ImageAnalysisSignature, - sub_lm=MagicMock(), - max_iterations=1, - sbx_pool=MagicMock(spec=SbxPool), - ) - - def test_sbx_pool_conflicts_with_custom_interpreter(self): - from predict_rlm import SbxPool - - with pytest.raises(ValueError, match="interpreter.*sbx_pool"): - PredictRLM( - ImageAnalysisSignature, - sub_lm=MagicMock(), - max_iterations=1, - sandbox_backend="sbx", - interpreter=MagicMock(), - sbx_pool=MagicMock(spec=SbxPool), - ) - - def test_sbx_pool_leases_without_constructing_or_shutting_down_interpreter(self): - from contextlib import contextmanager - - from predict_rlm import SbxPool - - leased = MagicMock() - leased.shutdown = MagicMock() - pool = MagicMock(spec=SbxPool) - - @contextmanager - def lease(**kwargs): - pool.lease_kwargs = kwargs - yield leased + assert result.answer == "done" + assert prompts == [1, 3] + assert [step.output for step in result.trace.steps] == [ + "confirm attempt 1", + "checked output", + "confirm attempt 3", + "FINAL: {'answer': 'done'}", + ] - pool.lease.side_effect = lease + @pytest.mark.asyncio + async def test_async_submit_confirmation_matches_sync_path(self): rlm = PredictRLM( ImageAnalysisSignature, sub_lm=MagicMock(), - max_iterations=1, - sandbox_backend="sbx", - sbx_pool=pool, - debug=True, + max_iterations=3, + submit_confirmation=lambda _context: "async confirm", ) - execution_tools = {"predict": MagicMock()} - - with patch("predict_rlm.backends.sbx.SbxBackend") as mock_sbx: - with rlm._interpreter_context(execution_tools=execution_tools) as repl: - assert repl is leased - pool.lease.assert_called_once() - assert pool.lease_kwargs["tools"] == execution_tools - assert pool.lease_kwargs["output_fields"] == rlm._get_output_fields_info() - assert pool.lease_kwargs["debug"] is True - assert pool.lease_kwargs["verbose"] is True - mock_sbx.assert_not_called() - leased.shutdown.assert_not_called() + result = await self._run_async( + rlm, + [ + self._prediction("SUBMIT(answer='done')"), + self._prediction("SUBMIT(answer='done')"), + ], + ) + assert result.answer == "done" + assert result.trace.status == "completed" + assert [step.output for step in result.trace.steps] == [ + "async confirm", + "FINAL: {'answer': 'done'}", + ] -class TestPredictRLMOutputDefaults: - def test_output_field_defaults_are_registered_for_submit(self): - rlm = PredictRLM(DefaultAnswerSignature, max_iterations=1) - assert rlm._get_output_fields_info() == [ - { - "name": "answer", - "has_default": True, - "default": None, - } - ] +class TestFatalExecutionErrors: + def test_sync_sandbox_fatal_error_propagates(self): + from predict_rlm.backends.base import SandboxFatalError + mock_lm = MagicMock() + rlm = PredictRLM(ImageAnalysisSignature, sub_lm=mock_lm, max_iterations=5) -class TestVerboseDebugLogging: - def test_verbose_streams_iteration_header_before_execute_and_output_after(self, capsys): - rlm = PredictRLM( - ImageAnalysisSignature, - sub_lm=MagicMock(), - max_iterations=5, - verbose=True, - ) + mock_repl = MagicMock() + mock_repl.execute = MagicMock(side_effect=SandboxFatalError("fatal")) mock_pred = MagicMock() mock_pred.reasoning = "thinking" - mock_pred.code = "print('model authored')\nprint('done')" + mock_pred.code = "print('hello')" rlm.generate_action = MagicMock(return_value=mock_pred) - seen: dict[str, str] = {} - - class Repl: - def execute(self, code, variables=None): - from predict_rlm._logging import ( - emit_trace_tool_call, - live_tool_call_logging_enabled, - ) - - seen["before_execute"] = capsys.readouterr().err - assert live_tool_call_logging_enabled() is True - emit_trace_tool_call("lookup", args=["needle"], kwargs={"limit": 1}) - return "visible output" - - with patch.object(rlm, "_process_execution_result", return_value=MagicMock()): + with pytest.raises(SandboxFatalError, match="fatal"): rlm._execute_iteration( - repl=Repl(), + repl=mock_repl, variables=[], history=[], iteration=0, @@ -423,55 +336,25 @@ def execute(self, code, variables=None): output_field_names=["answer"], ) - before_execute = seen["before_execute"] - after_execute = capsys.readouterr().err - before_text = _strip_ansi(before_execute) - after_text = _strip_ansi(after_execute) - - _assert_raw_verbose_output(before_execute) - _assert_raw_verbose_output(after_execute) - assert "\033[1;97mRLM turn \033[1;32m1\033[1;97m/5" in before_execute - assert "RLM turn 1/5" in before_text - assert "reasoning:" in before_text - assert "thinking" in before_text - assert "code:" in before_text - assert "model authored" in before_text - assert "output:" not in before_text - - assert "Tool: lookup(" in after_text - assert '"args": ["needle"]' in after_text - assert '"kwargs": {"limit": 1}' in after_text - assert "output:" in after_text - assert "visible output" in after_text - - def test_debug_logs_lifecycle_without_verbose_trace(self, capsys, caplog): - caplog.set_level(logging.DEBUG, logger="predict_rlm") - rlm = PredictRLM( - ImageAnalysisSignature, - sub_lm=MagicMock(), - max_iterations=5, - debug=True, - verbose=False, - ) - - mock_pred = MagicMock() - mock_pred.reasoning = "thinking" - mock_pred.code = "print('model authored')" - rlm.generate_action = MagicMock(return_value=mock_pred) + @pytest.mark.asyncio + async def test_sandbox_fatal_error_propagates(self): + from predict_rlm.backends.base import SandboxFatalError - seen: dict[str, str] = {} + mock_lm = MagicMock() + rlm = PredictRLM(ImageAnalysisSignature, sub_lm=mock_lm, max_iterations=5) - class Repl: - def execute(self, code, variables=None): - from predict_rlm._logging import live_tool_call_logging_enabled + mock_repl = MagicMock() + mock_repl.aexecute = AsyncMock(side_effect=SandboxFatalError("fatal")) - seen["before_execute"] = capsys.readouterr().err - assert live_tool_call_logging_enabled() is False - return "visible output" + mock_pred = MagicMock() + mock_pred.reasoning = "thinking" + mock_pred.code = "print('hello')" + rlm.generate_action = MagicMock() + rlm.generate_action.acall = AsyncMock(return_value=mock_pred) - with patch.object(rlm, "_process_execution_result", return_value=MagicMock()): - rlm._execute_iteration( - repl=Repl(), + with pytest.raises(SandboxFatalError, match="fatal"): + await rlm._aexecute_iteration( + repl=mock_repl, variables=[], history=[], iteration=0, @@ -479,2414 +362,32 @@ def execute(self, code, variables=None): output_field_names=["answer"], ) - stderr = seen["before_execute"] + capsys.readouterr().err - assert "RLM turn" not in stderr - assert "reasoning:" not in stderr - assert "code:" not in stderr - assert "output:" not in stderr - events = [record.getMessage().split()[0] for record in caplog.records] - assert "rlm.action_generation.start" in events - assert "rlm.action_generation.ok" in events - assert "rlm.execute.start" in events - assert "rlm.execute.ok" in events - - -class TestPredictTool: - """Tests that PredictRLM predict tool correctly runs DSPy signatures.""" - - @pytest.mark.asyncio - async def test_predict_returns_dict_response(self): - """predict tool runs DSPy Predict and returns dict output.""" - mock_lm = MagicMock() - rlm = PredictRLM(ImageAnalysisSignature, sub_lm=mock_lm, max_iterations=5) - - with patch("predict_rlm.predict_rlm.dspy.Predict") as mock_predict_class: - mock_predictor = MagicMock() - mock_prediction = MagicMock() - values = {"answer": "Paris"} - mock_prediction.keys.return_value = list(values.keys()) - mock_prediction.__getitem__ = MagicMock(side_effect=lambda k: values[k]) - mock_predictor.acall = AsyncMock(return_value=mock_prediction) - mock_predict_class.return_value = mock_predictor - - result = await rlm.tools["predict"].func( - "question -> answer", - question="What is the capital of France?", - ) - - assert isinstance(result, dict) - assert result == {"answer": "Paris"} - # Predict is called with a parsed Signature object - mock_predict_class.assert_called_once() - sig = mock_predict_class.call_args[0][0] - assert hasattr(sig, "input_fields") and "question" in sig.input_fields - mock_predictor.acall.assert_called_once_with( - question="What is the capital of France?" - ) - - @pytest.mark.asyncio - async def test_predict_with_multiple_outputs(self): - """predict correctly handles signatures with multiple outputs.""" - mock_lm = MagicMock() - rlm = PredictRLM(ImageAnalysisSignature, sub_lm=mock_lm, max_iterations=5) - - with patch("predict_rlm.predict_rlm.dspy.Predict") as mock_predict_class: - mock_predictor = MagicMock() - mock_prediction = MagicMock() - values = {"title": "Test Document", "summary": "A brief summary"} - mock_prediction.keys.return_value = list(values.keys()) - mock_prediction.__getitem__ = MagicMock(side_effect=lambda k: values[k]) - mock_predictor.acall = AsyncMock(return_value=mock_prediction) - mock_predict_class.return_value = mock_predictor - - result = await rlm.tools["predict"].func( - "text -> title, summary", - text="Some document content", - ) - - assert isinstance(result, dict) - assert result == {"title": "Test Document", "summary": "A brief summary"} - - @pytest.mark.asyncio - async def test_predict_with_instructions(self): - """predict passes instructions to create a Signature with instructions.""" - mock_lm = MagicMock() - rlm = PredictRLM(ImageAnalysisSignature, sub_lm=mock_lm, max_iterations=5) - - with patch("predict_rlm.predict_rlm.dspy.Predict") as mock_predict_class: - with patch("predict_rlm.predict_rlm.dspy.Signature") as mock_sig_class: - mock_predictor = MagicMock() - mock_prediction = MagicMock() - values = {"toxic": True} - mock_prediction.keys.return_value = list(values.keys()) - mock_prediction.__getitem__ = MagicMock(side_effect=lambda k: values[k]) - mock_predictor.acall = AsyncMock(return_value=mock_prediction) - mock_predict_class.return_value = mock_predictor - mock_sig_class.return_value = "mocked_signature" - - result = await rlm.tools["predict"].func( - "comment -> toxic: bool", - instructions="Mark as toxic if the comment includes insults.", - comment="You're an idiot!", - ) - - assert result == {"toxic": True} - mock_sig_class.assert_called_once_with( - "comment -> toxic: bool", "Mark as toxic if the comment includes insults." - ) - mock_predict_class.assert_called_once_with("mocked_signature") - - @pytest.mark.asyncio - async def test_predict_uses_sub_lm(self): - """predict uses the sub_lm when provided.""" - mock_lm = MagicMock() - rlm = PredictRLM(ImageAnalysisSignature, sub_lm=mock_lm, max_iterations=5) - - with patch("predict_rlm.predict_rlm.dspy.Predict") as mock_predict_class: - mock_predictor = MagicMock() - mock_prediction = MagicMock() - values = {"answer": "Test answer"} - mock_prediction.keys.return_value = list(values.keys()) - mock_prediction.__getitem__ = MagicMock(side_effect=lambda k: values[k]) - mock_predictor.acall = AsyncMock(return_value=mock_prediction) - mock_predict_class.return_value = mock_predictor - - result = await rlm.tools["predict"].func( - "question -> answer", - question="Test question", - ) - - assert result == {"answer": "Test answer"} - - @pytest.mark.asyncio - async def test_predict_error_when_no_lm(self): - """predict raises error when no LM is available.""" - rlm = PredictRLM(ImageAnalysisSignature, sub_lm=None, max_iterations=3) - - with dspy.context(lm=None): - with pytest.raises(RuntimeError, match="No LM available for predict"): - await rlm.tools["predict"].func("question -> answer", question="test") - - @pytest.mark.asyncio - async def test_predict_auto_wraps_images_with_type_hint(self): - """predict automatically wraps image URLs when field has dspy.Image type hint.""" - mock_lm = MagicMock() - rlm = PredictRLM(ImageAnalysisSignature, sub_lm=mock_lm, max_iterations=5) - - with patch("predict_rlm.predict_rlm.dspy.Predict") as mock_predict_class: - mock_predictor = MagicMock() - mock_prediction = MagicMock() - values = {"answer": "Extracted text"} - mock_prediction.keys.return_value = list(values.keys()) - mock_prediction.__getitem__ = MagicMock(side_effect=lambda k: values[k]) - mock_predictor.acall = AsyncMock(return_value=mock_prediction) - mock_predict_class.return_value = mock_predictor - - result = await rlm.tools["predict"].func( - "image: dspy.Image, question -> answer", - image="https://example.com/image.png", - question="What text is visible?", - ) - - assert result == {"answer": "Extracted text"} - # Verify the image was wrapped in dspy.Image - call_kwargs = mock_predictor.acall.call_args.kwargs - assert isinstance(call_kwargs["image"], dspy.Image) - assert call_kwargs["question"] == "What text is visible?" - - @pytest.mark.asyncio - async def test_predict_auto_wraps_base64_images_with_type_hint(self): - """predict automatically wraps base64 image data when field has dspy.Image type hint.""" - mock_lm = MagicMock() - rlm = PredictRLM(ImageAnalysisSignature, sub_lm=mock_lm, max_iterations=5) - - with patch("predict_rlm.predict_rlm.dspy.Predict") as mock_predict_class: - mock_predictor = MagicMock() - mock_prediction = MagicMock() - values = {"text": "OCR result"} - mock_prediction.keys.return_value = list(values.keys()) - mock_prediction.__getitem__ = MagicMock(side_effect=lambda k: values[k]) - mock_predictor.acall = AsyncMock(return_value=mock_prediction) - mock_predict_class.return_value = mock_predictor - - result = await rlm.tools["predict"].func( - "document: dspy.Image -> text", - document="data:image/png;base64,abc123...", - ) - - assert result == {"text": "OCR result"} - call_kwargs = mock_predictor.acall.call_args.kwargs - assert isinstance(call_kwargs["document"], dspy.Image) - - @pytest.mark.asyncio - async def test_predict_does_not_wrap_without_type_hint(self): - """predict does not wrap values for fields without dspy.Image type hint.""" - mock_lm = MagicMock() - rlm = PredictRLM(ImageAnalysisSignature, sub_lm=mock_lm, max_iterations=5) - - with patch("predict_rlm.predict_rlm.dspy.Predict") as mock_predict_class: - with patch("predict_rlm.predict_rlm.dspy.Image") as mock_image_class: - mock_predictor = MagicMock() - mock_prediction = MagicMock() - values = {"answer": "42"} - mock_prediction.keys.return_value = list(values.keys()) - mock_prediction.__getitem__ = MagicMock(side_effect=lambda k: values[k]) - mock_predictor.acall = AsyncMock(return_value=mock_prediction) - mock_predict_class.return_value = mock_predictor - - result = await rlm.tools["predict"].func( - "question -> answer", - question="https://example.com/some-url", - ) - - assert result == {"answer": "42"} - mock_image_class.assert_not_called() - mock_predictor.acall.assert_called_once_with( - question="https://example.com/some-url", - ) - @pytest.mark.asyncio - async def test_predict_uses_context_lm_captured_by_forward(self): - """predict uses context LM captured during forward() for thread-safe execution.""" - context_lm = MagicMock() - context_lm.name = "context_lm" - global_lm = MagicMock() - global_lm.name = "global_lm" - - rlm = PredictRLM(ImageAnalysisSignature, sub_lm=None, max_iterations=3) - - # Simulate forward() capturing the context LM - rlm._context_lm = context_lm - - with patch("predict_rlm.predict_rlm.dspy.Predict") as mock_predict_class: - with patch("predict_rlm.predict_rlm.dspy.context") as mock_context: - mock_predictor = MagicMock() - mock_prediction = MagicMock() - mock_prediction.keys.return_value = ["answer"] - mock_prediction.answer = "Test" - mock_predictor.acall = AsyncMock(return_value=mock_prediction) - mock_predict_class.return_value = mock_predictor - - _ = await rlm.tools["predict"].func( - "question -> answer", - question="Test?", - ) - - mock_context.assert_called_once_with(lm=context_lm) - - def test_forward_captures_and_clears_context_lm(self): - """forward() captures context LM before execution and clears it after.""" - context_lm = MagicMock() - rlm = PredictRLM(ImageAnalysisSignature, sub_lm=None, max_iterations=1) - - assert rlm._context_lm is None - - with patch.object(PredictRLM, "_forward_traced") as mock_traced: - mock_traced.return_value = dspy.Prediction(answer="Test") - - with dspy.context(lm=context_lm): - - def check_context_lm(file_plan, **kwargs): - assert rlm._context_lm is context_lm - return dspy.Prediction(answer="Test") - - mock_traced.side_effect = check_context_lm - - _ = rlm.forward(images=["img"], query="test?") - - assert rlm._context_lm is None - - @pytest.mark.asyncio - async def test_predict_auto_wraps_list_of_images_with_type_hint(self): - """predict automatically wraps list of image URLs when field has list[dspy.Image] type hint.""" - mock_lm = MagicMock() - rlm = PredictRLM(ImageAnalysisSignature, sub_lm=mock_lm, max_iterations=5) - - with patch("predict_rlm.predict_rlm.dspy.Predict") as mock_predict_class: - mock_predictor = MagicMock() - mock_prediction = MagicMock() - values = {"answer": "Analyzed 3 images"} - mock_prediction.keys.return_value = list(values.keys()) - mock_prediction.__getitem__ = MagicMock(side_effect=lambda k: values[k]) - mock_predictor.acall = AsyncMock(return_value=mock_prediction) - mock_predict_class.return_value = mock_predictor - - result = await rlm.tools["predict"].func( - "images: list[dspy.Image], question -> answer", - images=[ - "https://example.com/img1.png", - "https://example.com/img2.png", - "https://example.com/img3.png", - ], - question="What do these images show?", - ) - - assert result == {"answer": "Analyzed 3 images"} - # Predictor should receive list of wrapped dspy.Image instances - mock_predictor.acall.assert_called_once() - call_kwargs = mock_predictor.acall.call_args.kwargs - assert len(call_kwargs["images"]) == 3 - assert all(isinstance(img, dspy.Image) for img in call_kwargs["images"]) - assert call_kwargs["question"] == "What do these images show?" - - @pytest.mark.asyncio - async def test_predict_records_failed_subcall_in_trace_collector(self): - """predict failures are still recorded for structured tracing.""" - mock_lm = MagicMock() - mock_lm.model = "test-sub-lm" - mock_lm.history = [] - rlm = PredictRLM(ImageAnalysisSignature, sub_lm=mock_lm, max_iterations=5) - - init_predict_call_collector() - - with patch("predict_rlm.predict_rlm.dspy.Predict") as mock_predict_class: - mock_predictor = MagicMock() - mock_predictor.acall = AsyncMock(side_effect=RuntimeError("subcall boom")) - mock_predict_class.return_value = mock_predictor - - with pytest.raises(RuntimeError, match="subcall boom"): - await rlm.tools["predict"].func( - "question -> answer", - question="What is the capital of France?", - ) - - groups = drain_predict_calls() - assert len(groups) == 1 - assert groups[0].signature == "question -> answer" - assert len(groups[0].calls) == 1 - assert groups[0].calls[0].error == "subcall boom" - assert groups[0].calls[0].input == {"question": "What is the capital of France?"} - assert groups[0].calls[0].output == {} - - -class TestTypeContractEnforcement: - """Tests for type contract enforcement on predict outputs.""" - - @pytest.mark.asyncio - async def test_none_for_non_optional_list_raises(self): - """LM returning None for non-Optional list[X] field raises RuntimeError.""" - mock_lm = MagicMock() - rlm = PredictRLM(ImageAnalysisSignature, sub_lm=mock_lm, max_iterations=5) - - with patch("predict_rlm.predict_rlm.dspy.Predict") as mock_predict_class: - mock_predictor = MagicMock() - mock_prediction = MagicMock() - values = {"items": None} - mock_prediction.keys.return_value = list(values.keys()) - mock_prediction.__getitem__ = MagicMock(side_effect=lambda k: values[k]) - mock_predictor.acall = AsyncMock(return_value=mock_prediction) - mock_predict_class.return_value = mock_predictor - - with pytest.raises(RuntimeError, match="LM returned None for non-Optional"): - await rlm.tools["predict"].func( - "text: str -> items: list[str]", - text="some input", - ) - - @pytest.mark.asyncio - async def test_empty_list_passes_through_unchanged(self): - """Empty list [] is always valid — it means 'nothing found', distinct from None.""" - mock_lm = MagicMock() - rlm = PredictRLM(ImageAnalysisSignature, sub_lm=mock_lm, max_iterations=5) - - with patch("predict_rlm.predict_rlm.dspy.Predict") as mock_predict_class: - mock_predictor = MagicMock() - mock_prediction = MagicMock() - values = {"items": []} - mock_prediction.keys.return_value = list(values.keys()) - mock_prediction.__getitem__ = MagicMock(side_effect=lambda k: values[k]) - mock_predictor.acall = AsyncMock(return_value=mock_prediction) - mock_predict_class.return_value = mock_predictor - - result = await rlm.tools["predict"].func( - "text: str -> items: list[str]", - text="some input", - ) - assert result["items"] == [] - - @pytest.mark.asyncio - async def test_none_for_optional_str_passes_through(self): - """None for Optional[str] field passes through without error.""" - mock_lm = MagicMock() - rlm = PredictRLM(ImageAnalysisSignature, sub_lm=mock_lm, max_iterations=5) - - with patch("predict_rlm.predict_rlm.dspy.Predict") as mock_predict_class: - mock_predictor = MagicMock() - mock_prediction = MagicMock() - values = {"answer": None} - mock_prediction.keys.return_value = list(values.keys()) - mock_prediction.__getitem__ = MagicMock(side_effect=lambda k: values[k]) - mock_predictor.acall = AsyncMock(return_value=mock_prediction) - mock_predict_class.return_value = mock_predictor - - result = await rlm.tools["predict"].func( - "text: str -> answer: Optional[str]", - text="some input", - ) - assert result["answer"] is None - - @pytest.mark.asyncio - async def test_none_for_non_optional_str_raises(self): - """LM returning None for non-Optional str field raises RuntimeError.""" - mock_lm = MagicMock() - rlm = PredictRLM(ImageAnalysisSignature, sub_lm=mock_lm, max_iterations=5) - - with patch("predict_rlm.predict_rlm.dspy.Predict") as mock_predict_class: - mock_predictor = MagicMock() - mock_prediction = MagicMock() - values = {"answer": None} - mock_prediction.keys.return_value = list(values.keys()) - mock_prediction.__getitem__ = MagicMock(side_effect=lambda k: values[k]) - mock_predictor.acall = AsyncMock(return_value=mock_prediction) - mock_predict_class.return_value = mock_predictor - - with pytest.raises(RuntimeError, match="LM returned None for non-Optional"): - await rlm.tools["predict"].func( - "text: str -> answer: str", - text="some input", - ) - - @pytest.mark.asyncio - async def test_type_contract_failure_is_recorded_in_trace_collector(self): - """Post-call validation failures still record the predict attempt.""" - mock_lm = MagicMock() - mock_lm.model = "test-sub-lm" - mock_lm.history = [] - rlm = PredictRLM(ImageAnalysisSignature, sub_lm=mock_lm, max_iterations=5) - - init_predict_call_collector() - - with patch("predict_rlm.predict_rlm.dspy.Predict") as mock_predict_class: - mock_predictor = MagicMock() - mock_prediction = MagicMock() - values = {"answer": None} - mock_prediction.keys.return_value = list(values.keys()) - mock_prediction.__getitem__ = MagicMock(side_effect=lambda k: values[k]) - mock_predictor.acall = AsyncMock(return_value=mock_prediction) - mock_predict_class.return_value = mock_predictor - - with pytest.raises(RuntimeError, match="LM returned None for non-Optional"): - await rlm.tools["predict"].func( - "text: str -> answer: str", - text="some input", - ) - - groups = drain_predict_calls() - assert len(groups) == 1 - assert len(groups[0].calls) == 1 - assert "LM returned None for non-Optional" in groups[0].calls[0].error - assert groups[0].calls[0].input == {"text": "some input"} - - -class TestPredictRLMConfiguration: - """Tests for PredictRLM configuration options.""" - - def test_predict_always_exists(self): - """predict tool is always available (uses context LM if sub_lm not provided).""" - rlm = PredictRLM(ImageAnalysisSignature, sub_lm=None, max_iterations=5) - assert "predict" in rlm.tools - - def test_user_predict_not_overwritten(self): - """User-provided predict is not replaced.""" - mock_lm = MockLM() - - def user_predict(signature: str, **kwargs) -> dict: - return {"answer": "user implementation"} - - rlm = PredictRLM( - ImageAnalysisSignature, - sub_lm=mock_lm, - tools={"predict": user_predict}, - max_iterations=5, - ) - - result = rlm.tools["predict"].func("question -> answer", question="test") - assert result == {"answer": "user implementation"} - assert len(mock_lm.calls) == 0 - - def test_instructions_reference_predict_not_llm_query(self): - """PredictRLM instructions mention predict, not llm_query or sub_lm_query.""" - rlm = PredictRLM(ImageAnalysisSignature, sub_lm=None, max_iterations=5) - instructions = str(rlm.generate_action.signature.instructions) - - assert "predict" in instructions - assert "llm_query" not in instructions - assert "sub_lm_query" not in instructions - - def test_action_instructions_omit_timeout_guidance_by_default(self): - """Model-chosen execution timeouts are off by default: no field, no prompt.""" - rlm = PredictRLM(ImageAnalysisSignature, sub_lm=None, max_iterations=5) - action = rlm.generate_action.signature - instructions = " ".join(str(action.instructions).split()) - assert "execution_timeout_seconds" not in action.output_fields - assert "execution_timeout_seconds" not in instructions - assert "Execution timeouts" not in instructions - - def test_action_instructions_steer_iteration_execution_timeout(self): - """With model_execution_timeout=True the prompt explains when/how to use timeouts.""" - rlm = PredictRLM( - ImageAnalysisSignature, - sub_lm=None, - max_iterations=5, - model_execution_timeout=True, - ) - instructions = " ".join(str(rlm.generate_action.signature.instructions).split()) - - assert "execution_timeout_seconds" in instructions - assert "For every iteration" in instructions - assert "Use `null` for ordinary short, safe blocks" in instructions - assert "loops" in instructions - assert "scans over many files/items" in instructions - assert "network or tool fanout" in instructions - assert "batch `predict()` calls" in instructions - assert "tests/subprocesses" in instructions - assert "~1-5 seconds" in instructions - assert "~10-60 seconds" in instructions - assert "stdout/stderr printed before the timeout are preserved" in instructions - assert "next iteration can continue" in instructions - assert "store important partial results in variables" in instructions - assert "Strongly avoid long blocking calls" not in instructions - assert "Some operations may not be interruptible" not in instructions - assert "outside the Python event loop" not in instructions - assert "run bounded probes first" not in instructions - assert "scale up only after you understand the cost" not in instructions - assert "Prefer staged verification" not in instructions - assert "large SQLite queries" not in instructions - assert "`fetchall()` over unknown result sizes" not in instructions - assert "Do not run a full baseline query" not in instructions - - def test_allowed_domains_passed_to_interpreter(self): - """PredictRLM passes allowed_domains to interpreter.""" - rlm = PredictRLM( - ImageAnalysisSignature, - sub_lm=None, - max_iterations=5, - allowed_domains=["api.example.com"], - ) - assert rlm._allowed_domains == ["api.example.com"] - - -class TestMainLMParameter: - """Tests for the lm parameter on PredictRLM.""" - - def test_lm_as_dspy_lm_instance(self): - """Passing a dspy.LM instance copies it (fresh history) so - concurrent PredictRLM instances don't share mutable state. - """ - mock_lm = MagicMock(spec=dspy.LM) - mock_lm.copy.return_value = MagicMock(spec=dspy.LM) - rlm = PredictRLM(ImageAnalysisSignature, lm=mock_lm, max_iterations=1) - # rlm._lm is the COPY, not the original - assert rlm._lm is mock_lm.copy.return_value - mock_lm.copy.assert_called_once() - - def test_lm_as_string_creates_dspy_lm(self): - """Passing a model string creates a dspy.LM instance directly - (no copy needed — it's already a fresh instance). - """ - with patch("predict_rlm.predict_rlm.dspy.LM") as mock_lm_class: - mock_lm_class.return_value = MagicMock() - rlm = PredictRLM(ImageAnalysisSignature, lm="openai/gpt-4o", max_iterations=1) - mock_lm_class.assert_any_call("openai/gpt-4o", cache=False) - assert rlm._lm is mock_lm_class.return_value - - def test_lm_none_by_default(self): - """lm defaults to None (uses context LM).""" - rlm = PredictRLM(ImageAnalysisSignature, max_iterations=1) - assert rlm._lm is None - - def test_forward_uses_lm_as_context(self): - """forward() wraps execution in dspy.context(lm=...) using the - per-RLM copy (not the original passed in) so the context LM has - an isolated history. - """ - mock_lm = MagicMock() - mock_lm_copy = MagicMock() - mock_lm.copy.return_value = mock_lm_copy - rlm = PredictRLM(ImageAnalysisSignature, lm=mock_lm, max_iterations=1) - - with patch.object(PredictRLM, "_forward_traced") as mock_traced: - mock_traced.return_value = dspy.Prediction(answer="Test") - - captured_lm = None - - def capture_context(file_plan, **kwargs): - nonlocal captured_lm - captured_lm = dspy.settings.lm - return dspy.Prediction(answer="Test") - - mock_traced.side_effect = capture_context - rlm.forward(images=["img"], query="test?") - # The context LM is our PRIVATE copy, not the original - assert captured_lm is mock_lm_copy - - def test_forward_without_lm_uses_external_context(self): - """forward() without lm uses whatever is in dspy.context.""" - external_lm = MagicMock() - rlm = PredictRLM(ImageAnalysisSignature, lm=None, max_iterations=1) - - with patch.object(PredictRLM, "_forward_traced") as mock_traced: - mock_traced.return_value = dspy.Prediction(answer="Test") - - captured_lm = None - - def capture_context(file_plan, **kwargs): - nonlocal captured_lm - captured_lm = dspy.settings.lm - return dspy.Prediction(answer="Test") - - mock_traced.side_effect = capture_context - - with dspy.context(lm=external_lm): - rlm.forward(images=["img"], query="test?") - - assert captured_lm is external_lm - - def test_forward_clears_context_lm_after_execution(self): - """forward() clears _context_lm after execution when lm is provided.""" - mock_lm = MagicMock() - rlm = PredictRLM(ImageAnalysisSignature, lm=mock_lm, max_iterations=1) - - with patch.object(PredictRLM, "_forward_traced") as mock_traced: - mock_traced.return_value = dspy.Prediction(answer="Test") - rlm.forward(images=["img"], query="test?") - - assert rlm._context_lm is None - - def test_forward_clears_context_lm_on_error(self): - """forward() clears _context_lm even if execution raises.""" - mock_lm = MagicMock() - rlm = PredictRLM(ImageAnalysisSignature, lm=mock_lm, max_iterations=1) - - with patch.object(PredictRLM, "_forward_traced") as mock_traced: - mock_traced.side_effect = RuntimeError("boom") - - with pytest.raises(RuntimeError): - rlm.forward(images=["img"], query="test?") - - assert rlm._context_lm is None - - def test_lm_and_sub_lm_both_accepted(self): - """Both lm and sub_lm can be provided together; each is copied - for per-instance history isolation. - """ - mock_lm = MagicMock(spec=dspy.LM) - mock_sub_lm = MagicMock(spec=dspy.LM) - mock_lm.copy.return_value = MagicMock(spec=dspy.LM) - mock_sub_lm.copy.return_value = MagicMock(spec=dspy.LM) - rlm = PredictRLM( - ImageAnalysisSignature, - lm=mock_lm, - sub_lm=mock_sub_lm, - max_iterations=1, - ) - assert rlm._lm is mock_lm.copy.return_value - assert rlm._sub_lm is mock_sub_lm.copy.return_value - - @pytest.mark.asyncio - async def test_aforward_uses_lm_as_context(self): - """aforward() wraps execution in dspy.context(lm=...) using the - per-RLM copy (isolated history from the caller's original). - """ - mock_lm = MagicMock() - mock_lm_copy = MagicMock() - mock_lm.copy.return_value = mock_lm_copy - rlm = PredictRLM(ImageAnalysisSignature, lm=mock_lm, max_iterations=1) - - with patch.object(PredictRLM, "_aforward_traced") as mock_traced: - captured_lm = None - - async def capture_context(file_plan, **kwargs): - nonlocal captured_lm - captured_lm = dspy.settings.lm - return dspy.Prediction(answer="Test") - - mock_traced.side_effect = capture_context - await rlm.aforward(images=["img"], query="test?") - assert captured_lm is mock_lm_copy - - @pytest.mark.asyncio - async def test_aforward_clears_context_lm_after_execution(self): - """aforward() clears _context_lm after execution when lm is provided.""" - mock_lm = MagicMock() - rlm = PredictRLM(ImageAnalysisSignature, lm=mock_lm, max_iterations=1) - - with patch.object(PredictRLM, "_aforward_traced") as mock_traced: - mock_traced.return_value = dspy.Prediction(answer="Test") - await rlm.aforward(images=["img"], query="test?") - - assert rlm._context_lm is None - - @pytest.mark.asyncio - async def test_aforward_without_lm_uses_external_context(self): - """aforward() without lm uses whatever is in dspy.context.""" - external_lm = MagicMock() - rlm = PredictRLM(ImageAnalysisSignature, lm=None, max_iterations=1) - - with patch.object(PredictRLM, "_aforward_traced") as mock_traced: - captured_lm = None - - async def capture_context(file_plan, **kwargs): - nonlocal captured_lm - captured_lm = dspy.settings.lm - return dspy.Prediction(answer="Test") - - mock_traced.side_effect = capture_context - - with dspy.context(lm=external_lm): - await rlm.aforward(images=["img"], query="test?") - - assert captured_lm is external_lm - - -class TestTracedErrorHandling: - """Tests for trace data when traced execution fails.""" - - def test_forward_traced_error_attaches_error_trace(self): - mock_lm = MagicMock() - rlm = PredictRLM(ImageAnalysisSignature, sub_lm=mock_lm, max_iterations=2) - - repl = MagicMock() - context = MagicMock() - context.__enter__.return_value = repl - context.__exit__.return_value = False - - with ( - patch.object(PredictRLM, "_interpreter_context", return_value=context), - patch.object(PredictRLM, "_prepare_execution_tools", return_value={}), - patch.object(PredictRLM, "_build_variables", return_value={}), - patch.object(rlm, "_execute_iteration", side_effect=RuntimeError("boom")), - ): - with pytest.raises(RuntimeError, match="boom") as exc_info: - rlm._forward_traced(None, images=["img"], query="q") - - exc = exc_info.value - assert exc.trace.status == "error" - assert exc.trace.iterations == 0 - assert exc.trace.steps == [] - - def test_forward_traced_error_preserves_steps_before_failure(self): - mock_lm = MagicMock() - rlm = PredictRLM(ImageAnalysisSignature, sub_lm=mock_lm, max_iterations=2) - - repl = MagicMock() - context = MagicMock() - context.__enter__.return_value = repl - context.__exit__.return_value = False - - def side_effect(repl_obj, _variables, _history, iteration, _input_args, _output_fields): - if iteration == 0: - return REPLHistory( - entries=[ - REPLEntry(reasoning="first", code="print('ok')", output="ok"), - ] - ) - raise RuntimeError("boom") - - with ( - patch.object(PredictRLM, "_interpreter_context", return_value=context), - patch.object(PredictRLM, "_prepare_execution_tools", return_value={}), - patch.object(PredictRLM, "_build_variables", return_value={}), - patch.object(rlm, "_execute_iteration", side_effect=side_effect), - ): - with pytest.raises(RuntimeError, match="boom") as exc_info: - rlm._forward_traced(None, images=["img"], query="q") - - exc = exc_info.value - assert exc.trace.status == "error" - assert exc.trace.iterations == 1 - assert len(exc.trace.steps) == 1 - assert exc.trace.steps[0].iteration == 1 - assert exc.trace.steps[0].reasoning == "first" - - @pytest.mark.asyncio - async def test_aforward_traced_error_attaches_error_trace(self): - mock_lm = MagicMock() - rlm = PredictRLM(ImageAnalysisSignature, sub_lm=mock_lm, max_iterations=2) - - repl = MagicMock() - context = MagicMock() - context.__enter__.return_value = repl - context.__exit__.return_value = False - - with ( - patch.object(PredictRLM, "_interpreter_context", return_value=context), - patch.object(PredictRLM, "_prepare_execution_tools", return_value={}), - patch.object(PredictRLM, "_build_variables", return_value={}), - patch.object( - rlm, "_aexecute_iteration", new=AsyncMock(side_effect=RuntimeError("boom")) - ), - ): - with pytest.raises(RuntimeError, match="boom") as exc_info: - await rlm._aforward_traced(None, images=["img"], query="q") - - exc = exc_info.value - assert exc.trace.status == "error" - assert exc.trace.iterations == 0 - assert exc.trace.steps == [] - - -class TestModelsFromSchema: - """Tests for _models_from_schema function that reconstructs Pydantic models.""" - - def test_simple_model_from_schema(self): - """Simple model with basic fields is reconstructed correctly.""" - - class TaskItem(BaseModel): - category: str - title: str - - schema = TaskItem.model_json_schema() - models = _models_from_schema(schema) - - assert "TaskItem" in models - Model = models["TaskItem"] - - # Verify field names and types - assert set(Model.model_fields.keys()) == {"category", "title"} - - # Test instantiation - instance = Model(category="Test", title="My Task") - assert instance.category == "Test" - assert instance.title == "My Task" - - def test_model_with_optional_fields(self): - """Model with optional fields is reconstructed correctly.""" - from typing import Optional - - class Item(BaseModel): - name: str - description: Optional[str] = None - - schema = Item.model_json_schema() - models = _models_from_schema(schema) - - Model = models["Item"] - - # Test with optional field omitted - instance1 = Model(name="Widget") - assert instance1.name == "Widget" - assert instance1.description is None - - # Test with optional field provided - instance2 = Model(name="Widget", description="A useful widget") - assert instance2.description == "A useful widget" - - def test_model_with_list_fields(self): - """Model with list fields is reconstructed correctly.""" - from typing import List - - class Tags(BaseModel): - items: List[str] - counts: List[int] - - schema = Tags.model_json_schema() - models = _models_from_schema(schema) - - Model = models["Tags"] - instance = Model(items=["a", "b"], counts=[1, 2, 3]) - assert instance.items == ["a", "b"] - assert instance.counts == [1, 2, 3] - - def test_defaulted_fields_stay_non_nullable(self): - """Not-required fields must round-trip as omittable, NOT nullable. - - A field with a default (`mandatory: bool = True`) or a default_factory - (`tags: list[str]`) is "not required" in JSON Schema, but it is not - nullable. The reconstruction must keep the declared type (so the LM is told - the field is e.g. a list, never null) and preserve/restore the default -- - otherwise the LM emits null and the user's original model rejects it. - Regression for the predict() schema round-trip bug. - """ - from typing import List, Optional - - class Item(BaseModel): - name: str # required - note: Optional[str] = None # genuinely nullable - mandatory: bool = True # default value - tags: list[str] = Field(default_factory=list) # default_factory - - Model = _models_from_schema(Item.model_json_schema())["Item"] - fields = Model.model_fields - - # Genuinely-nullable field stays Optional; defaulted ones do NOT become Optional. - assert fields["note"].annotation == Optional[str] - assert fields["mandatory"].annotation is bool - assert fields["tags"].annotation == List[str] - - # Defaults are preserved / restored. - assert Model(name="x").mandatory is True - assert Model(name="x").tags == [] - - # The reconstructed model REJECTS null for the non-nullable defaulted fields - # (i.e. the LM will not be told null is acceptable). - with pytest.raises(ValidationError): - Model(name="x", tags=None) - with pytest.raises(ValidationError): - Model(name="x", mandatory=None) - - def test_nested_model_from_schema(self): - """Nested models with $defs are reconstructed correctly.""" - - class Address(BaseModel): - street: str - city: str - - class Person(BaseModel): - name: str - address: Address - - schema = Person.model_json_schema() - models = _models_from_schema(schema) - - # Both models should be created - assert "Person" in models - assert "Address" in models - - # Test instantiation with nested data - PersonModel = models["Person"] - AddressModel = models["Address"] - - addr = AddressModel(street="123 Main St", city="NYC") - person = PersonModel(name="Alice", address=addr) - assert person.name == "Alice" - assert person.address.street == "123 Main St" - - def test_deeply_nested_model(self): - """Deeply nested models are reconstructed correctly.""" - - class Country(BaseModel): - name: str - code: str - - class Address(BaseModel): - street: str - country: Country - - class Person(BaseModel): - name: str - address: Address - - schema = Person.model_json_schema() - models = _models_from_schema(schema) - - assert "Person" in models - assert "Address" in models - assert "Country" in models - - def test_model_with_list_of_nested_models(self): - """Model with list of nested models is reconstructed correctly.""" - from typing import List - - class LineItem(BaseModel): - product: str - quantity: int - - class Order(BaseModel): - order_id: str - items: List[LineItem] - - schema = Order.model_json_schema() - models = _models_from_schema(schema) - - assert "Order" in models - assert "LineItem" in models - - OrderModel = models["Order"] - LineItemModel = models["LineItem"] - - items = [ - LineItemModel(product="Widget", quantity=2), - LineItemModel(product="Gadget", quantity=1), - ] - order = OrderModel(order_id="ORD-123", items=items) - assert len(order.items) == 2 - assert order.items[0].product == "Widget" - - def test_model_with_all_primitive_types(self): - """Model with all supported primitive types is reconstructed.""" - - class AllTypes(BaseModel): - text: str - number: int - decimal: float - flag: bool - - schema = AllTypes.model_json_schema() - models = _models_from_schema(schema) - - Model = models["AllTypes"] - instance = Model(text="hello", number=42, decimal=3.14, flag=True) - assert instance.text == "hello" - assert instance.number == 42 - assert instance.decimal == 3.14 - assert instance.flag is True - - def test_enum_to_literal(self): - """Enum types in JSON schema are converted to Literal.""" - schema = { - "title": "Priority", - "properties": { - "level": {"enum": ["p1", "p2", "p3", "p4"], "type": "string"}, - }, - "required": ["level"], - } - models = _models_from_schema(schema) - Model = models["Priority"] - instance = Model(level="p1") - assert instance.level == "p1" - - def test_type_array_shorthand_for_optional(self): - """Type-array shorthand {"type": ["string", "null"]} → Optional[str].""" - schema = { - "title": "Item", - "properties": { - "name": {"type": "string"}, - "note": {"type": ["string", "null"]}, - }, - "required": ["name"], - } - models = _models_from_schema(schema) - Model = models["Item"] - instance = Model(name="test", note=None) - assert instance.name == "test" - assert instance.note is None - - def test_schema_without_title_gets_fallback(self): - """Schema without a title key uses the fallback name.""" - schema = { - "properties": {"x": {"type": "integer"}}, - "required": ["x"], - } - models = _models_from_schema(schema) - assert "RootModel" in models - - @pytest.mark.asyncio - async def test_predict_with_pydantic_schemas(self): - """predict tool uses pydantic_schemas to create custom_types.""" - mock_lm = MagicMock() - rlm = PredictRLM(ImageAnalysisSignature, sub_lm=mock_lm, max_iterations=5) - - # Create a schema that would come from sandbox - class TaskItem(BaseModel): - category: str - title: str - - pydantic_schemas = {"TaskItem": TaskItem.model_json_schema()} - - with patch("predict_rlm.predict_rlm.dspy.Predict") as mock_predict_class: - with patch("predict_rlm.predict_rlm.dspy.Signature") as mock_sig_class: - mock_predictor = MagicMock() - mock_prediction = MagicMock() - mock_prediction.keys.return_value = ["tasks"] - mock_prediction.tasks = [{"category": "Test", "title": "Task 1"}] - mock_predictor.acall = AsyncMock(return_value=mock_prediction) - mock_predict_class.return_value = mock_predictor - mock_sig_class.return_value = "mocked_signature" - - _ = await rlm.tools["predict"].func( - "text: str -> tasks: list[TaskItem]", - pydantic_schemas=pydantic_schemas, - text="test input", - ) - - # Verify Signature was called with custom_types - assert mock_sig_class.call_count == 1 - call_args = mock_sig_class.call_args - assert "custom_types" in call_args.kwargs - custom_types = call_args.kwargs["custom_types"] - assert "TaskItem" in custom_types - - @pytest.mark.asyncio - async def test_predict_without_pydantic_schemas_no_custom_types(self): - """predict without pydantic_schemas parses signature without custom_types.""" - mock_lm = MagicMock() - rlm = PredictRLM(ImageAnalysisSignature, sub_lm=mock_lm, max_iterations=5) - - with patch("predict_rlm.predict_rlm.dspy.Predict") as mock_predict_class: - mock_predictor = MagicMock() - mock_prediction = MagicMock() - mock_prediction.keys.return_value = ["answer"] - mock_prediction.answer = "Test" - mock_predictor.acall = AsyncMock(return_value=mock_prediction) - mock_predict_class.return_value = mock_predictor - - _ = await rlm.tools["predict"].func( - "question -> answer", - question="What is 2+2?", - ) - - # Predict should be called once with a parsed Signature - mock_predict_class.assert_called_once() - call_args = mock_predict_class.call_args - sig = call_args[0][0] - # Check it's a Signature object with expected fields - assert hasattr(sig, "input_fields") - assert hasattr(sig, "output_fields") - assert "question" in sig.input_fields - assert "answer" in sig.output_fields - - @pytest.mark.asyncio - async def test_predict_handles_items_field_without_collision(self): - """predict returns correct value when output field is named 'items'. - - Regression test: using getattr(prediction, 'items') returns the .items() - method instead of the field value. The fix uses prediction['items'] via - __getitem__ which bypasses method lookup. - """ - mock_lm = MagicMock() - rlm = PredictRLM(ImageAnalysisSignature, sub_lm=mock_lm, max_iterations=5) - - with patch("predict_rlm.predict_rlm.dspy.Predict") as mock_predict_class: - mock_predictor = MagicMock() - mock_prediction = MagicMock() - mock_prediction.keys.return_value = ["items"] - # Set up __getitem__ to return the actual value - expected_items = [{"title": "Task 1"}, {"title": "Task 2"}] - mock_prediction.__getitem__ = MagicMock(return_value=expected_items) - mock_predictor.acall = AsyncMock(return_value=mock_prediction) - mock_predict_class.return_value = mock_predictor - - result = await rlm.tools["predict"].func( - "page: dspy.Image -> items: list[dict]", - page="https://example.com/page.png", - ) - - assert isinstance(result, dict) - assert "items" in result - # Should return the list, not {} or [] from method collision - assert result["items"] == expected_items - mock_prediction.__getitem__.assert_called_once_with("items") - - -class TestAnnotationHelpers: - """Tests for _unwrap_optional, _image_field_info, _allows_none, _is_list_output. - - These are inner functions of _create_predict_tool, so we test them - indirectly through predict() behavior. - """ - - @pytest.mark.asyncio - async def test_optional_image_none_passes_through(self): - """Optional[dspy.Image] field with None value passes through without wrapping.""" - mock_lm = MagicMock() - rlm = PredictRLM(ImageAnalysisSignature, sub_lm=mock_lm, max_iterations=5) - - with patch("predict_rlm.predict_rlm.dspy.Predict") as mock_predict_class: - mock_predictor = MagicMock() - mock_prediction = MagicMock() - values = {"answer": "no image"} - mock_prediction.keys.return_value = list(values.keys()) - mock_prediction.__getitem__ = MagicMock(side_effect=lambda k: values[k]) - mock_predictor.acall = AsyncMock(return_value=mock_prediction) - mock_predict_class.return_value = mock_predictor - - result = await rlm.tools["predict"].func( - "image: Optional[dspy.Image], question -> answer", - image=None, - question="Any image?", - ) - - assert result == {"answer": "no image"} - call_kwargs = mock_predictor.acall.call_args.kwargs - assert call_kwargs["image"] is None - - @pytest.mark.asyncio - async def test_optional_list_image_wraps_correctly(self): - """Optional[list[dspy.Image]] wraps URLs as dspy.Image when list is provided.""" - mock_lm = MagicMock() - rlm = PredictRLM(ImageAnalysisSignature, sub_lm=mock_lm, max_iterations=5) - - with patch("predict_rlm.predict_rlm.dspy.Predict") as mock_predict_class: - mock_predictor = MagicMock() - mock_prediction = MagicMock() - values = {"answer": "found images"} - mock_prediction.keys.return_value = list(values.keys()) - mock_prediction.__getitem__ = MagicMock(side_effect=lambda k: values[k]) - mock_predictor.acall = AsyncMock(return_value=mock_prediction) - mock_predict_class.return_value = mock_predictor - - result = await rlm.tools["predict"].func( - "images: Optional[list[dspy.Image]], question -> answer", - images=["https://example.com/a.png", "https://example.com/b.png"], - question="Describe these", - ) - - assert result == {"answer": "found images"} - call_kwargs = mock_predictor.acall.call_args.kwargs - assert len(call_kwargs["images"]) == 2 - assert all(isinstance(img, dspy.Image) for img in call_kwargs["images"]) - - @pytest.mark.asyncio - async def test_none_for_optional_list_output_passes_through(self): - """None for Optional[list[str]] output passes through (allows_none is True).""" - mock_lm = MagicMock() - rlm = PredictRLM(ImageAnalysisSignature, sub_lm=mock_lm, max_iterations=5) - - with patch("predict_rlm.predict_rlm.dspy.Predict") as mock_predict_class: - mock_predictor = MagicMock() - mock_prediction = MagicMock() - values = {"items": None} - mock_prediction.keys.return_value = list(values.keys()) - mock_prediction.__getitem__ = MagicMock(side_effect=lambda k: values[k]) - mock_predictor.acall = AsyncMock(return_value=mock_prediction) - mock_predict_class.return_value = mock_predictor - - result = await rlm.tools["predict"].func( - "text: str -> items: Optional[list[str]]", - text="some input", - ) - assert result["items"] is None - - @pytest.mark.asyncio - async def test_is_list_output_detects_list_type(self): - """list[str] (non-Optional) output field: None raises RuntimeError (is_list + not allows_none).""" - mock_lm = MagicMock() - rlm = PredictRLM(ImageAnalysisSignature, sub_lm=mock_lm, max_iterations=5) - - with patch("predict_rlm.predict_rlm.dspy.Predict") as mock_predict_class: - mock_predictor = MagicMock() - mock_prediction = MagicMock() - values = {"tags": None} - mock_prediction.keys.return_value = list(values.keys()) - mock_prediction.__getitem__ = MagicMock(side_effect=lambda k: values[k]) - mock_predictor.acall = AsyncMock(return_value=mock_prediction) - mock_predict_class.return_value = mock_predictor - - with pytest.raises(RuntimeError, match="LM returned None for non-Optional"): - await rlm.tools["predict"].func( - "text: str -> tags: list[str]", - text="some input", - ) - - @pytest.mark.asyncio - async def test_non_optional_non_list_str_allows_none_false(self): - """Plain str output (not Optional): None raises RuntimeError.""" - mock_lm = MagicMock() - rlm = PredictRLM(ImageAnalysisSignature, sub_lm=mock_lm, max_iterations=5) - - with patch("predict_rlm.predict_rlm.dspy.Predict") as mock_predict_class: - mock_predictor = MagicMock() - mock_prediction = MagicMock() - values = {"name": None} - mock_prediction.keys.return_value = list(values.keys()) - mock_prediction.__getitem__ = MagicMock(side_effect=lambda k: values[k]) - mock_predictor.acall = AsyncMock(return_value=mock_prediction) - mock_predict_class.return_value = mock_predictor - - with pytest.raises(RuntimeError, match="LM returned None for non-Optional"): - await rlm.tools["predict"].func( - "text: str -> name: str", - text="some input", - ) - - -class TestSchemaTitleInjection: - """Tests for schema title injection when pydantic_schemas lack a 'title' key.""" - - @pytest.mark.asyncio - async def test_schema_without_title_injects_key_name(self): - """When pydantic_schemas has a schema missing 'title', the key name is injected.""" - mock_lm = MagicMock() - rlm = PredictRLM(ImageAnalysisSignature, sub_lm=mock_lm, max_iterations=5) - - schema_without_title = { - "properties": { - "description": {"type": "string"}, - "amount": {"type": "number"}, - }, - "required": ["description", "amount"], - } - - with patch("predict_rlm.predict_rlm.dspy.Predict") as mock_predict_class: - with patch("predict_rlm.predict_rlm.dspy.Signature") as mock_sig_class: - mock_predictor = MagicMock() - mock_prediction = MagicMock() - values = {"items": []} - mock_prediction.keys.return_value = list(values.keys()) - mock_prediction.__getitem__ = MagicMock(side_effect=lambda k: values[k]) - mock_predictor.acall = AsyncMock(return_value=mock_prediction) - mock_predict_class.return_value = mock_predictor - mock_sig_class.return_value = "mocked_signature" - - await rlm.tools["predict"].func( - "text: str -> items: list[LineItem]", - pydantic_schemas={"LineItem": schema_without_title}, - text="test", - ) - - call_args = mock_sig_class.call_args - custom_types = call_args.kwargs["custom_types"] - assert "LineItem" in custom_types - - @pytest.mark.asyncio - async def test_schema_with_title_preserved(self): - """When pydantic_schemas already has 'title', it is preserved (no overwrite).""" - mock_lm = MagicMock() - rlm = PredictRLM(ImageAnalysisSignature, sub_lm=mock_lm, max_iterations=5) - - schema_with_title = { - "title": "LineItem", - "properties": { - "description": {"type": "string"}, - }, - "required": ["description"], - } - - with patch("predict_rlm.predict_rlm.dspy.Predict") as mock_predict_class: - with patch("predict_rlm.predict_rlm.dspy.Signature") as mock_sig_class: - mock_predictor = MagicMock() - mock_prediction = MagicMock() - values = {"items": []} - mock_prediction.keys.return_value = list(values.keys()) - mock_prediction.__getitem__ = MagicMock(side_effect=lambda k: values[k]) - mock_predictor.acall = AsyncMock(return_value=mock_prediction) - mock_predict_class.return_value = mock_predictor - mock_sig_class.return_value = "mocked_signature" - - await rlm.tools["predict"].func( - "text: str -> items: list[LineItem]", - pydantic_schemas={"LineItem": schema_with_title}, - text="test", - ) - - call_args = mock_sig_class.call_args - custom_types = call_args.kwargs["custom_types"] - assert "LineItem" in custom_types - - -class TestUnresolvedTypesFallback: - """Tests for the fallback when signature has custom types that can't be resolved.""" - - @pytest.mark.asyncio - async def test_unresolved_custom_type_falls_back_to_string_signature(self, caplog): - """Unresolvable custom type in signature falls back to string signature with warning.""" - mock_lm = MagicMock() - rlm = PredictRLM(ImageAnalysisSignature, sub_lm=mock_lm, max_iterations=5) - - with patch("predict_rlm.predict_rlm.dspy.Predict") as mock_predict_class: - mock_predictor = MagicMock() - mock_prediction = MagicMock() - values = {"items": "raw string fallback"} - mock_prediction.keys.return_value = list(values.keys()) - mock_prediction.__getitem__ = MagicMock(side_effect=lambda k: values[k]) - mock_predictor.acall = AsyncMock(return_value=mock_prediction) - mock_predict_class.return_value = mock_predictor - - # Patch dspy.Signature to raise on unknown type, simulating a parse failure - original_sig = dspy.Signature - - def sig_side_effect(*args, **kwargs): - if kwargs.get("custom_types"): - return original_sig(*args, **kwargs) - sig_str = args[0] if args else "" - if "UnknownModel" in sig_str: - raise ValueError("Unknown name 'UnknownModel'") - return original_sig(*args, **kwargs) - - with patch("predict_rlm.predict_rlm.dspy.Signature", side_effect=sig_side_effect): - with caplog.at_level(logging.WARNING, logger="predict_rlm.predict_rlm"): - result = await rlm.tools["predict"].func( - "text: str -> items: list[UnknownModel]", - text="some input", - ) - - assert result == {"items": "raw string fallback"} - # Verify dspy.Predict was called with a string (the fallback) - sig_arg = mock_predict_class.call_args[0][0] - assert isinstance(sig_arg, str) - # Verify warning was logged about the fallback - assert any("UnknownModel" in r.message for r in caplog.records) - - @pytest.mark.asyncio - async def test_non_unknown_name_error_still_falls_back(self): - """Signature parse error without 'Unknown name' still falls back to string.""" - mock_lm = MagicMock() - rlm = PredictRLM(ImageAnalysisSignature, sub_lm=mock_lm, max_iterations=5) - - with patch("predict_rlm.predict_rlm.dspy.Predict") as mock_predict_class: - mock_predictor = MagicMock() - mock_prediction = MagicMock() - values = {"answer": "fallback"} - mock_prediction.keys.return_value = list(values.keys()) - mock_prediction.__getitem__ = MagicMock(side_effect=lambda k: values[k]) - mock_predictor.acall = AsyncMock(return_value=mock_prediction) - mock_predict_class.return_value = mock_predictor - - with patch( - "predict_rlm.predict_rlm.dspy.Signature", - side_effect=Exception("some parse error"), - ): - result = await rlm.tools["predict"].func( - "question -> answer", - question="test", - ) - - assert result == {"answer": "fallback"} - sig_arg = mock_predict_class.call_args[0][0] - assert isinstance(sig_arg, str) - - -class TestSubmitConfirmation: - """Tests for configurable submit confirmation in the main RLM loop.""" - - @staticmethod - def _prediction(code: str, reasoning: str = "thinking") -> dspy.Prediction: - return dspy.Prediction(reasoning=reasoning, code=code) - - def _run_sync( - self, - rlm: PredictRLM, - actions: list[dspy.Prediction], - repl: FakeSubmitRepl | None = None, - ) -> dspy.Prediction: - repl = repl or FakeSubmitRepl() - mock_lm = MagicMock() - mock_lm.history = [] - rlm.generate_action = MagicMock(side_effect=actions) - - with ( - dspy.context(lm=mock_lm), - patch.object(rlm, "_interpreter_context", return_value=FakeInterpreterContext(repl)), - ): - return rlm._forward_traced(None, images=["img"], query="Original task") - - async def _run_async( - self, - rlm: PredictRLM, - actions: list[dspy.Prediction], - repl: FakeSubmitRepl | None = None, - ) -> dspy.Prediction: - repl = repl or FakeSubmitRepl() - mock_lm = MagicMock() - mock_lm.history = [] - rlm.generate_action = MagicMock() - rlm.generate_action.acall = AsyncMock(side_effect=actions) - - with ( - dspy.context(lm=mock_lm), - patch.object(rlm, "_interpreter_context", return_value=FakeInterpreterContext(repl)), - ): - return await rlm._aforward_traced(None, images=["img"], query="Original task") - - def test_default_submit_completes_immediately(self): - rlm = PredictRLM(ImageAnalysisSignature, sub_lm=MagicMock(), max_iterations=3) - - result = self._run_sync( - rlm, - [self._prediction("SUBMIT(answer='done')")], - ) - - assert result.answer == "done" - assert result.trace.status == "completed" - assert len(result.trace.steps) == 1 - assert rlm.generate_action.call_count == 1 - - def test_max_output_chars_controls_traced_history_and_step_output(self): - rlm = PredictRLM( - ImageAnalysisSignature, - sub_lm=MagicMock(), - max_iterations=2, - max_output_chars=12, - ) - long_output = "abcdefghijklmnopqrstuvwxyz" - - class LongOutputRepl(FakeSubmitRepl): - def execute(self, code, variables=None, timeout=None): - self.executed.append(code) - if code.startswith("SUBMIT"): - return FinalOutput(self.final_payload) - return long_output - - result = self._run_sync( - rlm, - [ - self._prediction("print('long')"), - self._prediction("SUBMIT(answer='done')"), - ], - repl=LongOutputRepl(), - ) - - second_call_history = rlm.generate_action.call_args_list[1].kwargs["repl_history"] - formatted_history = second_call_history.format() - assert "Output (26 chars):" in formatted_history - assert "abcdefghijkl" not in formatted_history - assert "... (14 characters omitted) ..." in formatted_history - - first_step = result.trace.steps[0] - assert first_step.untruncated_output == long_output - assert first_step.output == "abcdefghijkl\n... (truncated to 12/26 chars)" - - def test_first_submit_prompts_and_second_submit_completes(self): - seen_contexts = [] - - def confirm(context): - seen_contexts.append(context) - return "Please verify the answer before final submit." - - rlm = PredictRLM( - ImageAnalysisSignature, - sub_lm=MagicMock(), - max_iterations=3, - submit_confirmation=confirm, - ) - - repl = FakeSubmitRepl() - result = self._run_sync( - rlm, - [ - self._prediction("SUBMIT(answer='done')", reasoning="first submit"), - self._prediction("SUBMIT(answer='done')", reasoning="second submit"), - ], - repl=repl, - ) - - assert result.answer == "done" - assert result.trace.status == "completed" - assert [step.output for step in result.trace.steps] == [ - "Please verify the answer before final submit.", - "FINAL: {'answer': 'done'}", - ] - assert rlm.generate_action.call_count == 2 - assert repl.deferred_submit_count == 1 - assert len(seen_contexts) == 1 - context = seen_contexts[0] - assert context.inputs == {"images": ["img"], "query": "Original task"} - assert context.output_field_names == ("answer",) - assert context.submitted_payload == {"answer": "done"} - assert context.prediction.answer == "done" - assert context.reasoning == "first submit" - assert context.code == "SUBMIT(answer='done')" - assert context.iteration == 1 - assert "FINAL" in context.latest_observation - assert len(context.history.entries) == 1 - - def test_confirmation_callback_can_skip_with_none_or_empty_string(self): - callbacks = [lambda _context: None, lambda _context: ""] - - for callback in callbacks: - rlm = PredictRLM( - ImageAnalysisSignature, - sub_lm=MagicMock(), - max_iterations=3, - submit_confirmation=callback, - ) - - result = self._run_sync( - rlm, - [self._prediction("SUBMIT(answer='done')")], - ) - - assert result.answer == "done" - assert result.trace.status == "completed" - assert len(result.trace.steps) == 1 - assert rlm.generate_action.call_count == 1 - - def test_non_submit_after_confirmation_clears_pending_confirmation(self): - prompts = [] - - def confirm(context): - prompts.append(context.iteration) - return f"confirm attempt {context.iteration}" - - rlm = PredictRLM( - ImageAnalysisSignature, - sub_lm=MagicMock(), - max_iterations=5, - submit_confirmation=confirm, - ) - - result = self._run_sync( - rlm, - [ - self._prediction("SUBMIT(answer='done')"), - self._prediction("print('checking')"), - self._prediction("SUBMIT(answer='done')"), - self._prediction("SUBMIT(answer='done')"), - ], - ) - - assert result.answer == "done" - assert prompts == [1, 3] - assert [step.output for step in result.trace.steps] == [ - "confirm attempt 1", - "checked output", - "confirm attempt 3", - "FINAL: {'answer': 'done'}", - ] - assert rlm.generate_action.call_count == 4 - - @pytest.mark.asyncio - async def test_async_submit_confirmation_matches_sync_path(self): - rlm = PredictRLM( - ImageAnalysisSignature, - sub_lm=MagicMock(), - max_iterations=3, - submit_confirmation=lambda _context: "async confirm", - ) - - result = await self._run_async( - rlm, - [ - self._prediction("SUBMIT(answer='done')"), - self._prediction("SUBMIT(answer='done')"), - ], - ) - - assert result.answer == "done" - assert result.trace.status == "completed" - assert [step.output for step in result.trace.steps] == [ - "async confirm", - "FINAL: {'answer': 'done'}", - ] - assert rlm.generate_action.acall.call_count == 2 - - -class TestExecuteIteration: - """Tests for _execute_iteration sync-path behavior.""" - - def test_action_lm_trace_metadata_stays_compact_with_prompt_cache_stats(self): - mock_lm = MagicMock() - mock_lm.history = [] - rlm = PredictRLM(ImageAnalysisSignature, sub_lm=mock_lm, max_iterations=5) - - mock_repl = MagicMock() - mock_repl.execute = MagicMock(return_value="output from execute") - - mock_pred = MagicMock() - mock_pred.reasoning = "thinking" - mock_pred.code = "print('hello')" - - def generate_action(**_kwargs): - mock_lm.history.append( - { - "usage": { - "prompt_tokens": 2000, - "completion_tokens": 100, - "prompt_tokens_details": {"cached_tokens": 1536}, - }, - "response": {"choices": [{"finish_reason": "stop"}]}, - } - ) - return mock_pred - - rlm.generate_action = MagicMock(side_effect=generate_action) - - with dspy.context(lm=mock_lm): - with patch.object(rlm, "_process_execution_result", return_value=MagicMock()): - rlm._execute_iteration( - repl=mock_repl, - variables=[], - history=[], - iteration=0, - input_args={}, - output_field_names=["answer"], - ) - - assert rlm._last_action_lm_metadata == LMFinishMetadata(finish_reason="stop") - - def test_accepts_repl_fence_in_sync_path(self): - mock_lm = MagicMock() - rlm = PredictRLM(ImageAnalysisSignature, sub_lm=mock_lm, max_iterations=5) - - mock_repl = MagicMock() - mock_repl.execute = MagicMock(return_value="output from execute") - - mock_pred = MagicMock() - mock_pred.reasoning = "thinking" - mock_pred.code = "```repl\nprint('hello')\n```" - rlm.generate_action = MagicMock(return_value=mock_pred) - - mock_result = MagicMock() - with patch.object(rlm, "_process_execution_result", return_value=mock_result): - result = rlm._execute_iteration( - repl=mock_repl, - variables=[], - history=[], - iteration=0, - input_args={}, - output_field_names=["answer"], - ) - - mock_repl.execute.assert_called_once_with("print('hello')", variables={}) - assert result is mock_result - - def test_verbose_streams_reasoning_and_code_before_sync_execute(self, capsys): - mock_lm = MagicMock() - rlm = PredictRLM( - ImageAnalysisSignature, - sub_lm=mock_lm, - max_iterations=5, - verbose=True, - ) - - mock_repl = MagicMock() - seen: dict[str, str] = {} - - def execute(code, variables=None): - seen["before_execute"] = capsys.readouterr().err - return "output from execute" - - mock_repl.execute = MagicMock(side_effect=execute) - - mock_pred = MagicMock() - mock_pred.reasoning = "thinking" - mock_pred.code = "```python\nprint('model authored')\n```" - rlm.generate_action = MagicMock(return_value=mock_pred) - - with patch.object(rlm, "_process_execution_result", return_value=MagicMock()): - rlm._execute_iteration( - repl=mock_repl, - variables=[], - history=[], - iteration=0, - input_args={"internal": "host value"}, - output_field_names=["answer"], - ) - - before_execute = seen["before_execute"] - after_execute = capsys.readouterr().err - before_text = _strip_ansi(before_execute) - after_text = _strip_ansi(after_execute) - _assert_raw_verbose_output(before_execute) - _assert_raw_verbose_output(after_execute) - assert "RLM turn 1/5" in before_text - assert "reasoning:" in before_text - assert "thinking" in before_text - assert "code:" in before_text - assert "model authored" in before_text - assert "output:" not in before_text - assert "output:" in after_text - assert "output from execute" in after_text - - def test_sync_sandbox_fatal_error_propagates(self): - from predict_rlm.backends.base import SandboxFatalError - - mock_lm = MagicMock() - rlm = PredictRLM(ImageAnalysisSignature, sub_lm=mock_lm, max_iterations=5) - - mock_repl = MagicMock() - mock_repl.execute = MagicMock(side_effect=SandboxFatalError("fatal")) - - mock_pred = MagicMock() - mock_pred.reasoning = "thinking" - mock_pred.code = "print('hello')" - rlm.generate_action = MagicMock(return_value=mock_pred) - - with pytest.raises(SandboxFatalError, match="fatal"): - rlm._execute_iteration( - repl=mock_repl, - variables=[], - history=[], - iteration=0, - input_args={}, - output_field_names=["answer"], - ) - - def test_failed_iteration_preserves_partial_output_before_error(self): - from predict_rlm.backends.base import SandboxExecutionError - - mock_lm = MagicMock() - rlm = PredictRLM(ImageAnalysisSignature, sub_lm=mock_lm, max_iterations=5) - - mock_repl = MagicMock() - mock_repl.execute = MagicMock( - side_effect=SandboxExecutionError( - "ValueError: bad", - partial_output="before failure\n", - ) - ) - - mock_pred = MagicMock() - mock_pred.reasoning = "thinking" - mock_pred.code = "print('before failure')\nraise ValueError('bad')" - rlm.generate_action = MagicMock(return_value=mock_pred) - - captured: dict[str, str] = {} - - def process_result(*args): - captured["result"] = args[2] if len(args) == 5 else args[1] - return MagicMock() - - with patch.object(rlm, "_process_execution_result", side_effect=process_result): - rlm._execute_iteration( - repl=mock_repl, - variables=[], - history=[], - iteration=0, - input_args={}, - output_field_names=["answer"], - ) - - assert captured["result"] == "before failure\n[Error] ValueError: bad" - - -class TestPredictRLMTelemetry: - """Focused generated-code telemetry tests without real LM calls.""" - - def test_interpreter_construction_receives_current_telemetry_context(self): - sink = ListTelemetrySink() - telemetry_context = TelemetryContext(sink=sink, trace_id="trace_case_interpreter") - rlm = PredictRLM( - ImageAnalysisSignature, - sub_lm=MagicMock(), - telemetry_context=telemetry_context, - ) - created_kwargs = {} - - class FakeJspiBackend: - def __init__(self, **kwargs): - created_kwargs.update(kwargs) - - def shutdown(self): - created_kwargs["shutdown_called"] = True - - with patch("predict_rlm.predict_rlm.JspiBackend", FakeJspiBackend): - rlm._begin_telemetry_execution() - try: - with rlm._interpreter_context(execution_tools={}) as repl: - assert isinstance(repl, FakeJspiBackend) - finally: - rlm._clear_telemetry_execution() - - assert created_kwargs["telemetry_context"] is telemetry_context - assert created_kwargs["shutdown_called"] is True - - def test_generated_code_event_uses_safe_payload(self): - sink = ListTelemetrySink() - telemetry_context = TelemetryContext(sink=sink, trace_id="trace_case_1") - rlm = PredictRLM( - ImageAnalysisSignature, - sub_lm=MagicMock(), - max_iterations=5, - telemetry_context=telemetry_context, - ) - - mock_repl = MagicMock() - mock_repl.execute = MagicMock(return_value="output") - mock_pred = MagicMock() - mock_pred.reasoning = "I will inspect the data." - mock_pred.code = "```python\nsecret_code = 41 + 1\nprint(secret_code)\n```" - rlm.generate_action = MagicMock(return_value=mock_pred) - - with patch.object(rlm, "_process_execution_result", return_value=MagicMock()): - rlm._begin_telemetry_execution() - try: - rlm._execute_iteration( - repl=mock_repl, - variables=[], - history=[], - iteration=0, - input_args={}, - output_field_names=["answer"], - ) - finally: - rlm._clear_telemetry_execution() - - names = [record["name"] for record in sink.records] - assert names == [ - "rlm.action_generation.start", - "rlm.action_generation.ok", - "rlm.iteration.generated_code", - ] - event = next( - record - for record in sink.records - if record["name"] == "rlm.iteration.generated_code" - ) - attrs = event["attributes"] - code = "secret_code = 41 + 1\nprint(secret_code)" - assert attrs["iteration"] == 1 - assert attrs["has_code"] is True - assert attrs["code_chars"] == len(code) - assert ( - attrs["code_sha256"] == "sha256_" + hashlib.sha256(code.encode("utf-8")).hexdigest() - ) - assert attrs["reasoning_chars"] == len(mock_pred.reasoning) - assert "predict_rlm.predictor_id" in attrs - assert "reasoning" not in attrs - assert code not in str(attrs) - assert "secret_code" not in str(attrs) - - def test_invalid_action_output_emits_parse_error_classifiable_as_no_code(self): - sink = ListTelemetrySink() - telemetry_context = TelemetryContext(sink=sink, trace_id="trace_case_2") - rlm = PredictRLM( - ImageAnalysisSignature, - sub_lm=MagicMock(), - max_iterations=5, - telemetry_context=telemetry_context, - ) - - mock_pred = MagicMock() - mock_pred.reasoning = "thinking" - mock_pred.code = "" - rlm.generate_action = MagicMock(return_value=mock_pred) - - rlm._begin_telemetry_execution() - try: - with pytest.raises(RuntimeError, match="invalid code"): - rlm._execute_iteration( - repl=MagicMock(), - variables=[], - history=[], - iteration=0, - input_args={}, - output_field_names=["answer"], - ) - finally: - rlm._clear_telemetry_execution() - - event = next( - record - for record in sink.records - if record["name"] == "rlm.action_generation.parse_error" - ) - attrs = event["attributes"] - assert attrs["iteration"] == 1 - assert attrs["has_code"] is False - assert attrs["code_chars"] == 0 - assert attrs["failure.class"] == "model_no_code_generated" - assert classify_failure(None, [event]) == "model_no_code_generated" - - def test_action_generation_exception_emits_unknown_parse_error_evidence(self): - sink = ListTelemetrySink() - telemetry_context = TelemetryContext(sink=sink, trace_id="trace_case_3") - rlm = PredictRLM( - ImageAnalysisSignature, - sub_lm=MagicMock(), - max_iterations=5, - telemetry_context=telemetry_context, - ) - rlm.generate_action = MagicMock(side_effect=ConnectionError("lm unavailable")) - - rlm._begin_telemetry_execution() - try: - with pytest.raises(ConnectionError, match="lm unavailable"): - rlm._execute_iteration( - repl=MagicMock(), - variables=[], - history=[], - iteration=0, - input_args={}, - output_field_names=["answer"], - ) - finally: - rlm._clear_telemetry_execution() - - event = next( - record - for record in sink.records - if record["name"] == "rlm.action_generation.parse_error" - ) - attrs = event["attributes"] - assert attrs["failure.class"] == "unknown" - assert attrs["error.type"] == "ConnectionError" - assert attrs["has_code"] is False - - def test_error_trace_gets_compact_telemetry_ref(self): - sink = ListTelemetrySink() - telemetry_context = TelemetryContext(sink=sink, trace_id="trace_case_4") - rlm = PredictRLM( - ImageAnalysisSignature, - sub_lm=MagicMock(), - max_iterations=2, - telemetry_context=telemetry_context, - ) - - repl = MagicMock() - context = MagicMock() - context.__enter__.return_value = repl - context.__exit__.return_value = False - - with ( - patch.object(PredictRLM, "_interpreter_context", return_value=context), - patch.object(PredictRLM, "_prepare_execution_tools", return_value={}), - patch.object(PredictRLM, "_build_variables", return_value={}), - patch.object(rlm, "_execute_iteration", side_effect=RuntimeError("boom")), - ): - with pytest.raises(RuntimeError, match="boom") as exc_info: - rlm._forward_traced(None, images=["img"], query="q") - - ref = exc_info.value.trace.telemetry_ref - assert ref["trace_id"] == "trace_case_4" - assert ref["predictor_id"].startswith("prlm_") - - -class TestAexecuteIteration: - """Tests for _aexecute_iteration: async vs sync interpreter dispatch.""" - - @pytest.mark.asyncio - async def test_async_action_lm_trace_metadata_stays_compact_with_prompt_cache_stats(self): - mock_lm = MagicMock() - mock_lm.history = [] - rlm = PredictRLM(ImageAnalysisSignature, sub_lm=mock_lm, max_iterations=5) - - mock_repl = MagicMock() - mock_repl.aexecute = AsyncMock(return_value="output from aexecute") - - mock_pred = MagicMock() - mock_pred.reasoning = "thinking" - mock_pred.code = "print('hello')" - - async def generate_action(**_kwargs): - mock_lm.history.append( - { - "usage": { - "prompt_tokens": 2000, - "completion_tokens": 100, - "prompt_tokens_details": {"cached_tokens": 1536}, - }, - "response": {"choices": [{"finish_reason": "stop"}]}, - } - ) - return mock_pred - - rlm.generate_action = MagicMock() - rlm.generate_action.acall = AsyncMock(side_effect=generate_action) - - with dspy.context(lm=mock_lm): - with patch.object(rlm, "_process_execution_result", return_value=MagicMock()): - await rlm._aexecute_iteration( - repl=mock_repl, - variables=[], - history=[], - iteration=0, - input_args={}, - output_field_names=["answer"], - ) - - assert rlm._last_action_lm_metadata == LMFinishMetadata(finish_reason="stop") - - @pytest.mark.asyncio - async def test_uses_aexecute_when_available(self): - """_aexecute_iteration calls repl.aexecute() when it has the method.""" - mock_lm = MagicMock() - rlm = PredictRLM(ImageAnalysisSignature, sub_lm=mock_lm, max_iterations=5) - - mock_repl = MagicMock() - mock_repl.aexecute = AsyncMock(return_value="output from aexecute") - - mock_pred = MagicMock() - mock_pred.reasoning = "thinking" - mock_pred.code = "print('hello')" - - rlm.generate_action = MagicMock() - rlm.generate_action.acall = AsyncMock(return_value=mock_pred) - - mock_result = MagicMock() - with patch.object( - rlm, "_process_execution_result", return_value=mock_result - ) as mock_process: - with patch( - "predict_rlm.predict_rlm.strip_code_fences", return_value="print('hello')" - ): - result = await rlm._aexecute_iteration( - repl=mock_repl, - variables=[], - history=[], - iteration=0, - input_args={}, - output_field_names=["answer"], - ) - - mock_repl.aexecute.assert_called_once_with("print('hello')", variables={}) - assert result is mock_result - from predict_rlm.predict_rlm import _PARENT_TAKES_CODE - - if _PARENT_TAKES_CODE: - mock_process.assert_called_once_with( - mock_pred, "print('hello')", "output from aexecute", [], ["answer"] - ) - else: - mock_process.assert_called_once_with( - mock_pred, "output from aexecute", [], ["answer"] - ) - - @pytest.mark.asyncio - async def test_verbose_streams_reasoning_and_code_before_async_execute(self, capsys): - mock_lm = MagicMock() - rlm = PredictRLM( - ImageAnalysisSignature, - sub_lm=mock_lm, - max_iterations=5, - verbose=True, - ) - - mock_repl = MagicMock() - seen: dict[str, str] = {} - - async def aexecute(code, variables=None): - seen["before_execute"] = capsys.readouterr().err - return "output from aexecute" - - mock_repl.aexecute = AsyncMock(side_effect=aexecute) - - mock_pred = MagicMock() - mock_pred.reasoning = "thinking" - mock_pred.code = "```python\nprint('async model authored')\n```" - - rlm.generate_action = MagicMock() - rlm.generate_action.acall = AsyncMock(return_value=mock_pred) - - with patch.object(rlm, "_process_execution_result", return_value=MagicMock()): - await rlm._aexecute_iteration( - repl=mock_repl, - variables=[], - history=[], - iteration=0, - input_args={"internal": "host value"}, - output_field_names=["answer"], - ) - - before_execute = seen["before_execute"] - after_execute = capsys.readouterr().err - before_text = _strip_ansi(before_execute) - after_text = _strip_ansi(after_execute) - _assert_raw_verbose_output(before_execute) - _assert_raw_verbose_output(after_execute) - assert "RLM turn 1/5" in before_text - assert "reasoning:" in before_text - assert "thinking" in before_text - assert "code:" in before_text - assert "async model authored" in before_text - assert "output:" not in before_text - assert "output:" in after_text - assert "output from aexecute" in after_text - - @pytest.mark.asyncio - async def test_falls_back_to_execute_when_no_aexecute(self): - """_aexecute_iteration falls back to repl.execute() when aexecute is absent.""" - mock_lm = MagicMock() - rlm = PredictRLM(ImageAnalysisSignature, sub_lm=mock_lm, max_iterations=5) - - mock_repl = MagicMock(spec=[]) # empty spec = no attributes - mock_repl.execute = MagicMock(return_value="output from execute") - - mock_pred = MagicMock() - mock_pred.reasoning = "thinking" - mock_pred.code = "print('hi')" - - rlm.generate_action = MagicMock() - rlm.generate_action.acall = AsyncMock(return_value=mock_pred) - - mock_result = MagicMock() - with patch.object(rlm, "_process_execution_result", return_value=mock_result): - with patch( - "predict_rlm.predict_rlm.strip_code_fences", return_value="print('hi')" - ): - result = await rlm._aexecute_iteration( - repl=mock_repl, - variables=[], - history=[], - iteration=0, - input_args={}, - output_field_names=["answer"], - ) - - mock_repl.execute.assert_called_once_with("print('hi')", variables={}) - assert result is mock_result - - @pytest.mark.asyncio - async def test_catches_execution_exception(self): - """_aexecute_iteration catches exceptions from repl and formats as error.""" - mock_lm = MagicMock() - rlm = PredictRLM(ImageAnalysisSignature, sub_lm=mock_lm, max_iterations=5) - - mock_repl = MagicMock() - mock_repl.aexecute = AsyncMock(side_effect=RuntimeError("sandbox crashed")) - - mock_pred = MagicMock() - mock_pred.reasoning = "thinking" - mock_pred.code = "bad_code()" - - rlm.generate_action = MagicMock() - rlm.generate_action.acall = AsyncMock(return_value=mock_pred) - - mock_result = MagicMock() - with patch.object( - rlm, "_process_execution_result", return_value=mock_result - ) as mock_process: - with patch("predict_rlm.predict_rlm.strip_code_fences", return_value="bad_code()"): - await rlm._aexecute_iteration( - repl=mock_repl, - variables=[], - history=[], - iteration=0, - input_args={}, - output_field_names=["answer"], - ) - - from predict_rlm.predict_rlm import _PARENT_TAKES_CODE - - error_arg = mock_process.call_args[0][2 if _PARENT_TAKES_CODE else 1] - assert "[Error]" in error_arg - assert "sandbox crashed" in error_arg - - @pytest.mark.asyncio - async def test_sandbox_fatal_error_propagates(self): - from predict_rlm.backends.base import SandboxFatalError - - mock_lm = MagicMock() - rlm = PredictRLM(ImageAnalysisSignature, sub_lm=mock_lm, max_iterations=5) - - mock_repl = MagicMock() - mock_repl.aexecute = AsyncMock(side_effect=SandboxFatalError("fatal")) - - mock_pred = MagicMock() - mock_pred.reasoning = "thinking" - mock_pred.code = "print('hello')" - rlm.generate_action = MagicMock() - rlm.generate_action.acall = AsyncMock(return_value=mock_pred) - - with pytest.raises(SandboxFatalError, match="fatal"): - await rlm._aexecute_iteration( - repl=mock_repl, - variables=[], - history=[], - iteration=0, - input_args={}, - output_field_names=["answer"], - ) - - -class TestAforwardTracedUsage: - """Usage accounting regression tests. - - PredictRLM reads usage from its per-instance ``self._lm.history`` - (via ``usage_since``). DSPy's ``BaseLM._process_lm_response`` - populates each history entry with ``usage`` (prompt/completion - tokens) and ``cost`` (from ``_hidden_params["response_cost"]``). - Across iterations, history grows and ``usage_since(lm, 0)`` sums - everything in the run. - """ - - @pytest.mark.asyncio - async def test_usage_since_sums_history_across_iterations(self): - """Two iterations each append an entry to ``lm.history``; - ``usage_since(lm, 0)`` returns the combined tokens + cost, - reflecting both calls. - """ - from predict_rlm.trace import usage_since - - mock_lm = MagicMock() - mock_lm.history = [ - # Iteration 1: retry (small empty-ish call, still billed) - {"usage": {"prompt_tokens": 50, "completion_tokens": 5}, "cost": 0.0001}, - # Iteration 2: full call with real tokens - {"usage": {"prompt_tokens": 240, "completion_tokens": 80}, "cost": 0.003}, - ] - - u = usage_since(mock_lm, 0) - assert u.input_tokens == 290 - assert u.output_tokens == 85 - assert u.cost == pytest.approx(0.0031) - - def test_debug_lm_metadata_reports_openai_prompt_cache_hits(self): - mock_lm = MagicMock() - mock_lm.history = [ - { - "usage": { - "prompt_tokens": 2000, - "completion_tokens": 100, - "prompt_tokens_details": {"cached_tokens": 1536}, - }, - "response": {"choices": [{"finish_reason": "stop"}]}, - } - ] - metadata = lm_completion_metadata_since(mock_lm, 0) - rlm = PredictRLM.__new__(PredictRLM) - - attrs = rlm._debug_lm_metadata(metadata) - - assert attrs["lm_prompt_tokens"] == 2000 - assert attrs["lm_cached_prompt_tokens"] == 1536 - assert attrs["lm_prompt_cache_read_ratio"] == pytest.approx(1536 / 2000) - assert attrs["lm_output_tokens"] == 100 - - -class TestSkillsMergeIntoInit: - """Tests for skills merging into PredictRLM.__init__.""" - - def test_skills_merge_instructions(self): - """Skills instructions are merged into _skill_instructions.""" - skill = Skill( - name="test-skill", - instructions="Use the frobnicator for all extraction.", - ) - rlm = PredictRLM(ImageAnalysisSignature, sub_lm=None, max_iterations=1, skills=[skill]) - assert "frobnicator" in rlm._skill_instructions - assert "test-skill" in rlm._skill_instructions - - def test_skills_merge_packages(self): - """Skills packages are merged into _skill_packages.""" - skill = Skill( - name="pkg-skill", - packages=["pdfplumber", "pillow"], - ) - rlm = PredictRLM(ImageAnalysisSignature, sub_lm=None, max_iterations=1, skills=[skill]) - assert "pdfplumber" in rlm._skill_packages - assert "pillow" in rlm._skill_packages - - def test_skills_merge_modules(self): - """Skills modules are merged into _skill_modules.""" - skill = Skill( - name="mod-skill", - modules={"helpers": "/path/to/helpers.py"}, - ) - rlm = PredictRLM(ImageAnalysisSignature, sub_lm=None, max_iterations=1, skills=[skill]) - assert rlm._skill_modules == {"helpers": "/path/to/helpers.py"} - - def test_skills_merge_tools(self): - """Skills tools are accessible on the RLM alongside predict.""" - - def my_tool(x: str) -> str: - """A custom tool.""" - return x - - skill = Skill( - name="tool-skill", - tools={"my_tool": my_tool}, - ) - rlm = PredictRLM(ImageAnalysisSignature, sub_lm=None, max_iterations=1, skills=[skill]) - assert "my_tool" in rlm.tools - assert "predict" in rlm.tools - - def test_skill_tool_conflicts_with_user_tool_raises(self): - """Tool name conflict between a skill and the tools parameter raises ValueError.""" - - def my_tool(x: str) -> str: - """A tool.""" - return x - - skill = Skill( - name="conflict-skill", - tools={"my_tool": my_tool}, - ) - with pytest.raises(ValueError, match="Tool name conflict.*my_tool"): - PredictRLM( - ImageAnalysisSignature, - sub_lm=None, - max_iterations=1, - skills=[skill], - tools={"my_tool": my_tool}, - ) - - def test_multiple_skills_merge(self): - """Multiple skills have their instructions, packages, and tools merged.""" - - def tool_a() -> str: - """Tool A.""" - return "a" - - def tool_b() -> str: - """Tool B.""" - return "b" - - skill_a = Skill( - name="skill-a", - instructions="Use approach A.", - packages=["pkg-a"], - tools={"tool_a": tool_a}, - ) - skill_b = Skill( - name="skill-b", - instructions="Use approach B.", - packages=["pkg-b", "pkg-a"], # duplicate pkg-a - tools={"tool_b": tool_b}, - ) - rlm = PredictRLM( - ImageAnalysisSignature, - sub_lm=None, - max_iterations=1, - skills=[skill_a, skill_b], - ) - assert "approach A" in rlm._skill_instructions - assert "approach B" in rlm._skill_instructions - assert "pkg-a" in rlm._skill_packages - assert "pkg-b" in rlm._skill_packages - # Deduplicated packages - assert rlm._skill_packages.count("pkg-a") == 1 - assert "tool_a" in rlm.tools - assert "tool_b" in rlm.tools - assert "predict" in rlm.tools - - def test_no_skills_leaves_defaults_empty(self): - """Without skills, skill fields are empty.""" - rlm = PredictRLM(ImageAnalysisSignature, sub_lm=None, max_iterations=1) - assert rlm._skill_instructions == "" - assert rlm._skill_packages == [] - assert rlm._skill_modules == {} - - -class TestModelsFromSchemaEdgeCases: - """Tests for _models_from_schema edge cases: unknown $ref, anyOf with all nulls.""" - - def test_unknown_ref_falls_back_to_dict(self): - """$ref to a name not in $defs falls back to dict type.""" - schema = { - "title": "Container", - "properties": { - "data": {"$ref": "#/$defs/MissingModel"}, - }, - "required": ["data"], - } - models = _models_from_schema(schema) - Model = models["Container"] - # Should accept a dict since the ref fell back to dict type - instance = Model(data={"key": "value"}) - assert instance.data == {"key": "value"} - - def test_anyof_all_null_falls_back_to_optional_str(self): - """anyOf with only null types falls back to Optional[str].""" - schema = { - "title": "Weird", - "properties": { - "field": {"anyOf": [{"type": "null"}]}, - }, - "required": ["field"], - } - models = _models_from_schema(schema) - Model = models["Weird"] - instance = Model(field=None) - assert instance.field is None +def test_iteration_usage_is_not_charged_again_on_the_next_step(): + rlm = PredictRLM("q -> a", sub_lm=DummyLM([])) + rlm._last_action_lm_usage = TokenUsage(input_tokens=2000, output_tokens=100, cost=0.012) + groups = [ + PredictCallGroup( + signature="x -> y", + instructions=None, + model="dummy", + total_usage=TokenUsage(input_tokens=50, output_tokens=10, cost=0.001), + calls=[], + ), + PredictCallGroup( + signature="x -> z", + instructions=None, + model="dummy", + total_usage=TokenUsage(input_tokens=30, output_tokens=5, cost=0.0005), + calls=[], + ), + ] + usage = rlm._build_iteration_usage(groups) + assert usage.main == TokenUsage(input_tokens=2000, output_tokens=100, cost=0.012) + assert usage.sub.input_tokens == 80 + assert usage.sub.output_tokens == 15 + assert usage.sub.cost == pytest.approx(0.0015) + + next_usage = rlm._build_iteration_usage([]) + assert next_usage.main == TokenUsage() + assert next_usage.sub == TokenUsage() diff --git a/tests/test_pydantic_fixes.py b/tests/test_pydantic_fixes.py deleted file mode 100644 index cde05ef7..00000000 --- a/tests/test_pydantic_fixes.py +++ /dev/null @@ -1,973 +0,0 @@ -"""Tests for Pydantic serialization and sandbox detection fixes. - -These tests should fail without the fixes and pass with them. -""" - -import asyncio -import json -import re -import unittest -from typing import Literal, Optional -from unittest.mock import AsyncMock, MagicMock, patch - -import pytest -from pydantic import BaseModel, Field - - -def _run(coro): - """Run async predict call from sync test.""" - import nest_asyncio - - nest_asyncio.apply() - loop = asyncio.get_event_loop() - return loop.run_until_complete(coro) - - -class ExtractedItem(BaseModel): - """Test model with complex types including Literal.""" - - category: str = Field(description="Category") - title: str = Field(description="Title") - priority: Optional[Literal["urgent", "high", "medium", "low"]] = Field( - default=None, description="Priority" - ) - - -class TestPydanticSerialization(unittest.TestCase): - """Test that _to_serializable properly handles complex Pydantic types.""" - - def test_literal_field_serialization(self): - """Test that Pydantic models with Literal fields serialize without mode='python'.""" - # This test demonstrates the fix - without mode='python', Literal types don't serialize properly - - item = ExtractedItem( - category="TEST", - title="Test Task", - priority="high", # Literal field - ) - - # The OLD way (would have issues with Literal types) - def old_to_serializable(value): - if hasattr(value, "model_dump"): - return value.model_dump() # No mode='python' - return value - - # The NEW way (our fix) - def new_to_serializable(value): - if isinstance(value, BaseModel): - return value.model_dump(mode="python") # With mode='python' - return value - - # Both should produce dicts, but the new way handles Literal better - old_result = old_to_serializable(item) - new_result = new_to_serializable(item) - - # Both should be serializable to JSON - self.assertIsInstance(old_result, dict) - self.assertIsInstance(new_result, dict) - - # The new way should preserve the Literal value correctly - self.assertEqual(new_result["priority"], "high") - - # Both should be JSON serializable - json.dumps(old_result) - json.dumps(new_result) - - -class TestSandboxPydanticDetection(unittest.TestCase): - """Test that sandbox can detect Pydantic models defined in REPL.""" - - def test_signature_must_be_string(self): - """Test that _get_pydantic_schemas handles non-string signatures.""" - - # Simulate the pattern matching that happens in _get_pydantic_schemas - pattern = r":\s*(?:list\[|List\[|Optional\[)?([A-Z][A-Za-z0-9_]*)" - - # This would fail without our fix - sig_dict = {"signature": "test"} - - # OLD way - would crash with TypeError - with self.assertRaises(TypeError) as ctx: - list(re.finditer(pattern, sig_dict)) - self.assertIn("expected string or bytes-like object", str(ctx.exception)) - - # NEW way - convert to string first - sig_safe = str(sig_dict) if not isinstance(sig_dict, str) else sig_dict - matches = list(re.finditer(pattern, sig_safe)) - # Should not crash - self.assertIsInstance(matches, list) - - def test_pydantic_model_detection_in_different_scopes(self): - """Test that Pydantic models can be found in various scopes.""" - - # Simulate what _get_pydantic_schemas does - def find_model_in_scopes(name, test_globals=None, test_locals=None): - """Simplified version of our fixed _get_pydantic_schemas logic.""" - cls = None - - # Check provided scopes (for testing) - if test_globals and name in test_globals: - cls = test_globals[name] - elif test_locals and name in test_locals: - cls = test_locals[name] - - return cls - - # Test 1: Model in globals - test_globals = {"ExtractedItem": ExtractedItem} - found = find_model_in_scopes("ExtractedItem", test_globals=test_globals) - self.assertEqual(found, ExtractedItem) - - # Test 2: Model in locals (simulating REPL definition) - test_locals = {"ExtractedItem": ExtractedItem} - found = find_model_in_scopes("ExtractedItem", test_locals=test_locals) - self.assertEqual(found, ExtractedItem) - - # Test 3: Model not found - found = find_model_in_scopes("NonExistent", test_globals={}, test_locals={}) - self.assertIsNone(found) - - def test_schema_extraction_from_signature(self): - """Test that Pydantic schemas can be extracted from signatures.""" - - signature = "image: dspy.Image, doc_id: str -> items: list[ExtractedItem]" - - # Pattern to find type names - pattern = r":\s*(?:list\[|List\[|Optional\[)?([A-Z][A-Za-z0-9_]*)" - matches = re.finditer(pattern, signature) - - found_types = [] - for match in matches: - name = match.group(1) - # Skip DSPy built-ins - if name not in ("Image", "List", "Optional", "Dict"): - found_types.append(name) - - # Should find ExtractedItem - self.assertIn("ExtractedItem", found_types) - - # Now test schema extraction - # Simulate having ExtractedItem in scope - test_globals = {"ExtractedItem": ExtractedItem} - - schemas = {} - for type_name in found_types: - if type_name in test_globals: - cls = test_globals[type_name] - if hasattr(cls, "model_json_schema"): - schemas[type_name] = cls.model_json_schema() - - # Should have extracted the schema - self.assertIn("ExtractedItem", schemas) - self.assertIn("properties", schemas["ExtractedItem"]) - self.assertIn("category", schemas["ExtractedItem"]["properties"]) - - -class TestIntegration(unittest.TestCase): - """Test that the fixes work together in a realistic scenario.""" - - def test_predict_tool_with_pydantic_model(self): - """Test the full flow of using a Pydantic model in a predict signature.""" - - # This simulates what happens when the model uses ExtractedItem - signature = "page: dspy.Image -> items: list[ExtractedItem]" - - # Step 1: Ensure signature is string (our first fix) - if not isinstance(signature, str): - signature = str(signature) - - # Step 2: Extract Pydantic schemas (our second fix - checking multiple scopes) - pattern = r":\s*(?:list\[|List\[|Optional\[)?([A-Z][A-Za-z0-9_]*)" - schemas = {} - - # Simulate ExtractedItem being in the REPL's globals - repl_globals = {"ExtractedItem": ExtractedItem} - - for match in re.finditer(pattern, signature): - name = match.group(1) - if name == "ExtractedItem" and name in repl_globals: - cls = repl_globals[name] - if hasattr(cls, "model_json_schema"): - schemas[name] = cls.model_json_schema() - - # Should have found and extracted the schema - self.assertIn("ExtractedItem", schemas) - - # Step 3: Create payload with schemas - payload_dict = { - "args": [signature], - "kwargs": {"page": "image_url"}, - "pydantic_schemas": schemas, - } - - # Should be JSON serializable - json_payload = json.dumps(payload_dict) - self.assertIsInstance(json_payload, str) - - # Step 4: When predict returns ExtractedItem instances, they should serialize - result_items = [ExtractedItem(category="TASK", title="Test", priority="high")] - - # Our _to_serializable fix - def to_serializable(value): - if isinstance(value, BaseModel): - return value.model_dump(mode="python") - if isinstance(value, list): - return [to_serializable(item) for item in value] - return value - - serialized = to_serializable(result_items) - - # Should be JSON serializable - json_result = json.dumps(serialized) - self.assertIsInstance(json_result, str) - - # Should preserve the Literal value - parsed = json.loads(json_result) - self.assertEqual(parsed[0]["priority"], "high") - - -class TestListDictSerialization(unittest.TestCase): - """Test that list[dict] output type is properly serialized. - - This addresses the issue where using list[dict] as an output type - in predict() causes Pydantic serialization warnings about Message - and StreamingChoices objects from litellm. - """ - - def test_list_dict_serialization_in_to_serializable(self): - """Test that list[dict] values are properly handled by _to_serializable.""" - from typing import Any - - # Replicate the _to_serializable function from predict_rlm.py - def _to_serializable(value: Any) -> Any: - """Convert Pydantic models to dicts recursively.""" - if value is None: - return value - - if isinstance(value, BaseModel): - return value.model_dump(mode="python") - - if hasattr(value, "dict") and hasattr(value, "__fields__"): - return value.dict() - - if isinstance(value, list): - return [_to_serializable(item) for item in value] - - if isinstance(value, tuple): - return [_to_serializable(item) for item in value] - - if isinstance(value, dict): - return {k: _to_serializable(v) for k, v in value.items()} - - if isinstance(value, set): - return list(value) - - return value - - # Test: list of dicts with various value types - test_data = [ - {"category": "TASK", "title": "Test Task", "page": 1}, - {"category": "FORM", "title": "Another Task", "priority": "high"}, - ] - - result = _to_serializable(test_data) - - # Should be a list of dicts - self.assertIsInstance(result, list) - self.assertEqual(len(result), 2) - self.assertIsInstance(result[0], dict) - self.assertIsInstance(result[1], dict) - - # Should be JSON serializable - json_str = json.dumps(result) - self.assertIsInstance(json_str, str) - - # Parse back and verify - parsed = json.loads(json_str) - self.assertEqual(parsed[0]["category"], "TASK") - self.assertEqual(parsed[1]["priority"], "high") - - def test_list_dict_with_nested_pydantic_model(self): - """Test that list containing dicts with Pydantic models are serialized.""" - from typing import Any - - class NestedModel(BaseModel): - name: str - value: int - - def _to_serializable(value: Any) -> Any: - if value is None: - return value - if isinstance(value, BaseModel): - return value.model_dump(mode="python") - if isinstance(value, list): - return [_to_serializable(item) for item in value] - if isinstance(value, dict): - return {k: _to_serializable(v) for k, v in value.items()} - return value - - # A list of dicts where some values are Pydantic models - test_data = [ - {"category": "TASK", "nested": NestedModel(name="test", value=42)}, - ] - - result = _to_serializable(test_data) - - # Should have serialized the nested model - self.assertIsInstance(result[0]["nested"], dict) - self.assertEqual(result[0]["nested"]["name"], "test") - self.assertEqual(result[0]["nested"]["value"], 42) - - # Should be JSON serializable - json_str = json.dumps(result) - self.assertIsInstance(json_str, str) - - def test_list_dict_with_none_values(self): - """Test that list[dict] with None values is properly serialized.""" - - def _to_serializable(value): - if value is None: - return value - if isinstance(value, BaseModel): - return value.model_dump(mode="python") - if isinstance(value, list): - return [_to_serializable(item) for item in value] - if isinstance(value, dict): - return {k: _to_serializable(v) for k, v in value.items()} - return value - - # List of dicts with None values (common pattern in extraction results) - test_data = [ - { - "category": "TASK", - "title": "Test Task", - "priority": None, - "due_date": None, - }, - ] - - result = _to_serializable(test_data) - - # Should preserve None values - self.assertIsNone(result[0]["priority"]) - self.assertIsNone(result[0]["due_date"]) - - # Should be JSON serializable - json_str = json.dumps(result) - self.assertIn('"priority": null', json_str) - - -@pytest.mark.integration -class TestListDictInInterpreter(unittest.TestCase): - """Test list[dict] through the full interpreter pipeline.""" - - def test_predict_with_list_dict_output(self): - """Test that predict with list[dict] output works through the sandbox.""" - from predict_rlm.backends import JspiBackend - - # Track calls to predict - predict_calls = [] - - def mock_predict(signature: str, pydantic_schemas=None, **kwargs): - predict_calls.append( - {"signature": signature, "schemas": pydantic_schemas, "kwargs": kwargs} - ) - # Return a list of dicts (simulating DSPy output) - return { - "items": [ - {"category": "TASK", "title": "Task 1", "page": 0}, - {"category": "FORM", "title": "Form 1", "priority": "high"}, - ] - } - - interpreter = JspiBackend( - tools={"predict": mock_predict}, - preinstall_packages=False, - ) - try: - # This is the pattern from the user's trace that was causing issues - result = interpreter.execute(""" -import json - -result = await predict( - "context: str -> items: list[dict]", - instructions="Extract items", - context="test context" -) -print("Got items:", len(result["items"])) -print("First item:", result["items"][0]) - -# Verify we can serialize the result (this was failing) -serialized = json.dumps(result["items"]) -print("Serialized OK:", len(serialized), "chars") -""") - - # Check that it worked - assert "Got items: 2" in str(result) - assert "First item:" in str(result) - assert "Serialized OK:" in str(result) - - # Verify predict was called correctly - assert len(predict_calls) == 1 - assert predict_calls[0]["signature"] == "context: str -> items: list[dict]" - - finally: - interpreter.shutdown() - - def test_predict_with_list_dict_containing_none_values(self): - """Test that predict with list[dict] containing None values works correctly.""" - from predict_rlm.backends import JspiBackend - - def mock_predict(signature: str, **kwargs): - # Return list of dicts with None values (common in extraction results) - return { - "items": [ - { - "category": "TASK", - "title": "Task 1", - "priority": None, # Nullable field - "due_date": None, # Nullable field - "page": 0, - }, - { - "category": "FORM", - "title": "Form 1", - "priority": "high", - "due_date": "2026-01-15", - "page": 1, - }, - ] - } - - interpreter = JspiBackend( - tools={"predict": mock_predict}, - preinstall_packages=False, - ) - try: - result = interpreter.execute(""" -import json - -result = await predict( - "context: str -> items: list[dict]", - context="test context" -) - -# Check None values are preserved -first_item = result["items"][0] -print("First item priority:", first_item["priority"]) -print("First item due_date:", first_item["due_date"]) - -# Verify serialization works with None -serialized = json.dumps(result["items"]) -print("Serialized OK") - -# Verify None is preserved after serialization -parsed = json.loads(serialized) -print("Parsed priority is None:", parsed[0]["priority"] is None) -""") - - assert "First item priority: None" in str(result) - assert "First item due_date: None" in str(result) - assert "Serialized OK" in str(result) - assert "Parsed priority is None: True" in str(result) - - finally: - interpreter.shutdown() - - def test_predict_with_list_dict_complex_nested_values(self): - """Test that predict handles list[dict] with complex nested structures.""" - from predict_rlm.backends import JspiBackend - - def mock_predict(signature: str, **kwargs): - # Return list of dicts with nested structures - return { - "items": [ - { - "category": "TASK", - "title": "Task 1", - "metadata": { - "source": "RFP Section 5", - "references": ["page 1", "page 2"], - }, - "tags": ["urgent", "compliance"], - }, - ] - } - - interpreter = JspiBackend( - tools={"predict": mock_predict}, - preinstall_packages=False, - ) - try: - result = interpreter.execute(""" -import json - -result = await predict( - "context: str -> items: list[dict]", - context="test context" -) - -item = result["items"][0] -print("Metadata source:", item["metadata"]["source"]) -print("Tags:", item["tags"]) -print("References:", item["metadata"]["references"]) - -# Verify serialization works with nested structures -serialized = json.dumps(result["items"]) -print("Serialized length:", len(serialized)) -""") - - assert "Metadata source: RFP Section 5" in str(result) - assert "Tags: ['urgent', 'compliance']" in str(result) - assert "References: ['page 1', 'page 2']" in str(result) - assert "Serialized length:" in str(result) - - finally: - interpreter.shutdown() - - -class TestToSerializableEdgeCases(unittest.TestCase): - """Test edge cases in _to_serializable that could cause 'method is not JSON serializable'.""" - - def _get_to_serializable(self): - """Get the actual _to_serializable function from predict_rlm.""" - from typing import Any - - def _to_serializable(value: Any) -> Any: - """Convert Pydantic models and other objects to JSON-serializable dicts.""" - # Primitives pass through directly - if value is None or isinstance(value, (str, int, float, bool)): - return value - - if isinstance(value, BaseModel): - return value.model_dump(mode="python") - - if hasattr(value, "dict") and hasattr(value, "__fields__"): - return value.dict() - - # Dataclasses - convert to dict using asdict - if hasattr(value, "__dataclass_fields__"): - import dataclasses - - return {k: _to_serializable(v) for k, v in dataclasses.asdict(value).items()} - - if isinstance(value, list): - return [_to_serializable(item) for item in value] - - if isinstance(value, tuple): - return [_to_serializable(item) for item in value] - - if isinstance(value, dict): - return {k: _to_serializable(v) for k, v in value.items()} - - if isinstance(value, set): - return list(value) - - if hasattr(value, "__dict__") and type(value).__name__ == "Prediction": - return { - k: _to_serializable(v) - for k, v in value.__dict__.items() - if not k.startswith("_") - } - - # Fallback: try to convert to dict or string representation - if hasattr(value, "__dict__"): - try: - return { - k: _to_serializable(v) - for k, v in value.__dict__.items() - if not k.startswith("_") and not callable(v) - } - except Exception: - pass - - # Last resort: convert to string representation - return str(value) - - return _to_serializable - - def test_object_with_method_is_serializable(self): - """Test that objects with methods are now serializable (fixed).""" - _to_serializable = self._get_to_serializable() - - class ObjectWithMethod: - def __init__(self): - self.value = "test" - self.number = 42 - - def some_method(self): - return "result" - - obj = ObjectWithMethod() - result = _to_serializable(obj) - - # Now it should be a dict (methods excluded) - self.assertIsInstance(result, dict) - self.assertEqual(result["value"], "test") - self.assertEqual(result["number"], 42) - # Method should NOT be included - self.assertNotIn("some_method", result) - - # Should be JSON serializable - json_str = json.dumps(result) - self.assertIsInstance(json_str, str) - - def test_dataclass_is_serializable(self): - """Test that dataclasses are properly serialized (fixed).""" - from dataclasses import dataclass - - _to_serializable = self._get_to_serializable() - - @dataclass - class TaskItem: - category: str - title: str - priority: Optional[str] = None - - item = TaskItem(category="TASK", title="Test", priority="high") - result = _to_serializable(item) - - # Should be a dict - self.assertIsInstance(result, dict) - self.assertEqual(result["category"], "TASK") - self.assertEqual(result["title"], "Test") - self.assertEqual(result["priority"], "high") - - # Should be JSON serializable - json_str = json.dumps(result) - self.assertIsInstance(json_str, str) - - def test_nested_dataclass_is_serializable(self): - """Test that nested dataclasses are properly serialized.""" - from dataclasses import dataclass - - _to_serializable = self._get_to_serializable() - - @dataclass - class Address: - street: str - city: str - - @dataclass - class Person: - name: str - address: Address - - person = Person(name="Test", address=Address(street="123 Main", city="Boston")) - result = _to_serializable(person) - - # Should be nested dicts - self.assertIsInstance(result, dict) - self.assertEqual(result["name"], "Test") - self.assertIsInstance(result["address"], dict) - self.assertEqual(result["address"]["street"], "123 Main") - self.assertEqual(result["address"]["city"], "Boston") - - # Should be JSON serializable - json_str = json.dumps(result) - self.assertIsInstance(json_str, str) - - def test_list_of_dataclasses_is_serializable(self): - """Test that list of dataclasses is properly serialized.""" - from dataclasses import dataclass - - _to_serializable = self._get_to_serializable() - - @dataclass - class TaskItem: - category: str - title: str - - items = [ - TaskItem(category="TASK", title="Task 1"), - TaskItem(category="FORM", title="Form 1"), - ] - result = _to_serializable(items) - - # Should be a list of dicts - self.assertIsInstance(result, list) - self.assertEqual(len(result), 2) - self.assertIsInstance(result[0], dict) - self.assertEqual(result[0]["category"], "TASK") - - # Should be JSON serializable - json_str = json.dumps(result) - self.assertIsInstance(json_str, str) - - def test_unknown_object_converted_to_string(self): - """Test that completely unknown objects are converted to string.""" - _to_serializable = self._get_to_serializable() - - # An object without __dict__ (can't extract attributes) - class WeirdObject: - __slots__ = ["value"] - - def __init__(self): - self.value = "test" - - def __str__(self): - return "WeirdObject(value=test)" - - obj = WeirdObject() - result = _to_serializable(obj) - - # Should be converted to string - self.assertEqual(result, "WeirdObject(value=test)") - - # Should be JSON serializable - json_str = json.dumps(result) - self.assertIsInstance(json_str, str) - - -class TestPredictToolListDict: - """Test the actual predict tool from PredictRLM with list[dict] output.""" - - @pytest.mark.asyncio - async def test_predict_tool_with_list_dict_output(self): - """Test that the predict tool handles list[dict] correctly.""" - import warnings - - from predict_rlm import PredictRLM - - mock_lm = MagicMock() - rlm = PredictRLM("text -> answer", sub_lm=mock_lm, max_iterations=1) - predict_tool = rlm.tools["predict"].func - - with patch("predict_rlm.predict_rlm.dspy.Predict") as mock_predict_class: - mock_predictor = MagicMock() - mock_prediction = MagicMock() - mock_prediction.keys.return_value = ["items"] - items_list = [ - {"category": "TASK", "title": "Task 1", "page": 0, "priority": None}, - {"category": "FORM", "title": "Form 1", "page": 1, "priority": "high"}, - ] - mock_prediction.items = items_list - mock_prediction.__getitem__ = lambda self, key: getattr(self, key) - mock_predictor.acall = AsyncMock(return_value=mock_prediction) - mock_predict_class.return_value = mock_predictor - - with warnings.catch_warnings(record=True) as w: - warnings.simplefilter("always") - - result = await predict_tool( - "context: str -> items: list[dict]", - context="test context", - ) - - assert isinstance(result, dict) - assert "items" in result - assert isinstance(result["items"], list) - assert len(result["items"]) == 2 - assert result["items"][0]["category"] == "TASK" - assert result["items"][1]["priority"] == "high" - assert result["items"][0]["priority"] is None - - json_str = json.dumps(result) - assert isinstance(json_str, str) - - pydantic_warnings = [ - warning - for warning in w - if "PydanticSerializationUnexpectedValue" in str(warning.message) - ] - assert len(pydantic_warnings) == 0, ( - f"Got unexpected Pydantic warnings: {pydantic_warnings}" - ) - - @pytest.mark.asyncio - async def test_predict_tool_with_list_of_pydantic_models(self): - """Test that the predict tool handles list of Pydantic models correctly.""" - import warnings - - from predict_rlm import PredictRLM - - class TaskItem(BaseModel): - category: str - title: str - priority: Optional[str] = None - - mock_lm = MagicMock() - rlm = PredictRLM("text -> answer", sub_lm=mock_lm, max_iterations=1) - predict_tool = rlm.tools["predict"].func - - with patch("predict_rlm.predict_rlm.dspy.Predict") as mock_predict_class: - mock_predictor = MagicMock() - mock_prediction = MagicMock() - mock_prediction.keys.return_value = ["items"] - mock_prediction.items = [ - TaskItem(category="TASK", title="Task 1", priority=None), - TaskItem(category="FORM", title="Form 1", priority="high"), - ] - mock_prediction.__getitem__ = lambda self, key: getattr(self, key) - mock_predictor.acall = AsyncMock(return_value=mock_prediction) - mock_predict_class.return_value = mock_predictor - - with warnings.catch_warnings(record=True) as w: - warnings.simplefilter("always") - - result = await predict_tool( - "context: str -> items: list[TaskItem]", - context="test context", - pydantic_schemas={"TaskItem": TaskItem.model_json_schema()}, - ) - - assert isinstance(result, dict) - assert "items" in result - assert isinstance(result["items"], list) - assert isinstance(result["items"][0], dict) - assert result["items"][0]["category"] == "TASK" - - json_str = json.dumps(result) - assert isinstance(json_str, str) - - # Check for Pydantic warnings - pydantic_warnings = [ - warning - for warning in w - if "PydanticSerializationUnexpectedValue" in str(warning.message) - ] - assert len(pydantic_warnings) == 0 - - -@pytest.mark.integration -class TestSchemaExtractionFromCallStack(unittest.TestCase): - """Test that Pydantic schemas are extracted correctly from user's REPL scope.""" - - def test_schema_extraction_traverses_full_call_stack(self): - """Test that _get_pydantic_schemas traverses the entire call stack. - - This tests the fix for the issue where models defined in the REPL - weren't found because _get_pydantic_schemas only checked one frame back - instead of traversing the full call stack. - """ - import re - - # A model defined in this test's scope (simulates REPL-defined model) - class DeepStackModel(BaseModel): - name: str - value: int - - # Replicate the _get_pydantic_schemas logic with the FIX applied - def _get_pydantic_schemas_fixed(sig): - """Extract schemas by traversing the FULL call stack.""" - import inspect - - schemas = {} - pattern = r":\s*(?:list\[|List\[|Optional\[)?([A-Z][A-Za-z0-9_]*)" - for match in re.finditer(pattern, sig): - name = match.group(1) - if name in ("Image", "List", "Optional", "Dict", "Any", "Union"): - continue - - cls = None - # FIX: Traverse the FULL call stack, not just one frame back - frame = inspect.currentframe() - while frame: - if name in frame.f_globals: - cls = frame.f_globals[name] - break - if name in frame.f_locals: - cls = frame.f_locals[name] - break - frame = frame.f_back - - if cls and hasattr(cls, "model_json_schema"): - try: - schemas[name] = cls.model_json_schema() - except Exception: - pass - return schemas - - # Now call through multiple nested functions (simulating production call stack) - def level_3(sig): - return _get_pydantic_schemas_fixed(sig) - - def level_2(sig): - return level_3(sig) - - def level_1(sig): - return level_2(sig) - - # Call through deep stack - model is defined in THIS scope (like REPL) - schemas = level_1("context: str -> items: list[DeepStackModel]") - - # Should find the model even though it's many frames up - self.assertIn( - "DeepStackModel", - schemas, - f"DeepStackModel should be found in call stack. Got: {schemas}", - ) - self.assertIn("properties", schemas["DeepStackModel"]) - self.assertIn("name", schemas["DeepStackModel"]["properties"]) - - def test_schema_extraction_finds_model_in_caller_frame(self): - """Test that _get_pydantic_schemas can find models defined in caller's scope. - - This is a regression test for the issue where models defined in the REPL - weren't being found by _get_pydantic_schemas because it only looked at - immediate caller's globals, not the full call stack. - """ - from predict_rlm.backends import JspiBackend - - # Track what schemas are extracted and passed to predict - received_schemas = [] - - def mock_predict(signature: str, pydantic_schemas=None, **kwargs): - received_schemas.append( - { - "signature": signature, - "schemas": pydantic_schemas, - } - ) - return {"tasks": []} - - interpreter = JspiBackend( - tools={"predict": mock_predict}, - preinstall_packages=False, - ) - try: - # Define a Pydantic model in the REPL and use it in a signature - result = interpreter.execute(""" -from pydantic import BaseModel -from typing import Optional - -class TaskCandidate(BaseModel): - category: str - title: str - description: str - priority: Optional[str] = None - -# Call predict with the custom model in the signature -result = await predict( - "page: dspy.Image -> tasks: list[TaskCandidate]", - instructions="Extract tasks", - page="test_url" -) -print("Predict called successfully") -""") - - # Check that predict was called - assert "Predict called successfully" in str(result) - assert len(received_schemas) == 1 - - # Check that the schema was extracted and passed - schemas = received_schemas[0]["schemas"] - self.assertIsNotNone(schemas, "Schema should be extracted from REPL-defined model") - self.assertIn( - "TaskCandidate", schemas, f"TaskCandidate schema not found. Got: {schemas}" - ) - - # Verify the schema has the expected structure - tc_schema = schemas["TaskCandidate"] - self.assertIn("properties", tc_schema) - self.assertIn("category", tc_schema["properties"]) - self.assertIn("title", tc_schema["properties"]) - self.assertIn("description", tc_schema["properties"]) - self.assertIn("priority", tc_schema["properties"]) - - finally: - interpreter.shutdown() - - -if __name__ == "__main__": - unittest.main() diff --git a/tests/test_read_line_raw_hang.py b/tests/test_read_line_raw_hang.py deleted file mode 100644 index c67fc9f0..00000000 --- a/tests/test_read_line_raw_hang.py +++ /dev/null @@ -1,236 +0,0 @@ -"""RED-GREEN repro for the deno-stdout blocking-read hang. - -Background: - ``JspiBackend._read_with_timeout`` is called inline from - ``_send_request``, which is used by the synchronous health-check - path (``_health_check`` → ``_ensure_deno_process``). When that path - runs inside an asyncio coroutine (via ``aexecute`` → ``_aexecute_inner``), - a blocking read on the deno stdout fd freezes the **entire event loop**. - - The outer ``_read_with_timeout`` uses ``select.select()`` with a - timeout — but only to check *whether there is data to read*. Once - select says ready it calls ``_read_line_raw()``, which loops on - ``os.read(fd, 65536)`` until a newline arrives. If the first ``os.read`` - returns partial bytes without a newline, the next ``os.read`` iteration - has **no timeout** and blocks indefinitely waiting for more data. - - This caused a production stall on 2026-04-18: a gemini+medium eval - seized at 319/400 tasks with the main asyncio event loop parked in - ``os.read`` inside ``_read_line_raw``. ``kill -USR1`` via the - faulthandler committed that morning confirmed the exact frame. - -RED (pre-fix): partial bytes (no newline) followed by an unbounded wait - makes ``_read_with_timeout(timeout=0.2)`` hang forever. The test's - outer thread-join timeout of 2s catches the hang and fails the - assertion — proves the bug exists. - -GREEN (post-fix): ``_read_line_raw`` re-checks the deadline between - each ``os.read`` iteration and raises ``TimeoutError`` (or returns - ``None`` via ``_read_with_timeout``) when the budget is exhausted. -""" - -from __future__ import annotations - -import os -import threading -import time -import types - -import pytest - -from predict_rlm.backends import JspiBackend - - -def _make_interp_with_partial_pipe(): - """Build a JspiBackend pointed at a real OS pipe fd, with - some partial bytes (no newline) already sitting in the kernel - buffer. The write end is left open and silent, simulating a deno - process that wrote a partial line then stopped producing output - (crashed, hung on GC, blocked on its own I/O — doesn't matter). - """ - read_fd, write_fd = os.pipe() - # Write partial data (no '\n'). Writer stays open so read won't get EOF. - os.write(write_fd, b"partial-line-no-newline-here") - - interp = JspiBackend.__new__(JspiBackend) - interp._stdout_fd = read_fd - interp._read_buf = "" - interp.deno_process = types.SimpleNamespace( - stdout=types.SimpleNamespace(fileno=lambda: read_fd), - poll=lambda: None, - ) - return interp, write_fd - - -def test_read_with_timeout_does_not_hang_on_partial_line(): - """Hand ``_read_with_timeout`` an fd that emits partial bytes then - goes silent. It must return within the configured timeout (plus a - small safety margin) — not block forever. - - Run in a background thread so a hang in the code under test - doesn't freeze pytest itself. - """ - interp, write_fd = _make_interp_with_partial_pipe() - try: - result = {"returned": False, "value": None, "elapsed": 0.0} - - def _call(): - t0 = time.monotonic() - result["value"] = interp._read_with_timeout(timeout=0.2) - result["elapsed"] = time.monotonic() - t0 - result["returned"] = True - - t = threading.Thread(target=_call, daemon=True) - t.start() - # 2s safety net: if the bug is present the thread hangs here - # and we fail the assertion below with a clean message. - t.join(timeout=2.0) - - assert not t.is_alive(), ( - "JspiBackend._read_with_timeout hung on a partial-line " - "stdout — deno-stdout deadlock is present" - ) - assert result["returned"], "thread ended without returning a value" - # The caller configured timeout=0.2s, so the call should return - # within roughly that window plus a small safety margin. - assert result["elapsed"] < 1.0, ( - f"returned but took {result['elapsed']:.2f}s — longer than " - "the 0.2s caller-requested timeout, fix may be too loose" - ) - finally: - os.close(write_fd) - try: - os.close(interp._stdout_fd) - except OSError: - pass - - -def test_read_line_raw_directly_respects_timeout_when_given_one(): - """``_read_line_raw`` itself must accept a timeout (the fix's API - addition) and raise when exceeded. Narrower anchor than the - integration test above: fails if someone re-removes the timeout - parameter or short-circuits its check. - """ - interp, write_fd = _make_interp_with_partial_pipe() - try: - t0 = time.monotonic() - with pytest.raises((TimeoutError, OSError)): - interp._read_line_raw(timeout=0.1) - elapsed = time.monotonic() - t0 - assert elapsed < 0.5, ( - f"_read_line_raw(timeout=0.1) took {elapsed:.2f}s to fail — " - "timeout not being enforced tightly" - ) - finally: - os.close(write_fd) - try: - os.close(interp._stdout_fd) - except OSError: - pass - - -def test_send_request_does_not_hang_on_silent_deno(): - """The real-world trigger: ``_send_request`` calls - ``_read_with_timeout`` and must NOT pass ``timeout=None``. If it - does, a deno process that never replies freezes the caller (and, - when the caller is on the asyncio event loop, every sibling - coroutine with it). - - We mock the deno process's stdin (absorb writes silently) and stdout - (real pipe that emits partial bytes then stalls). ``_send_request`` - must raise within a bounded window — the previous hang behavior - would block forever and fail the 3s thread-join safety net. - """ - import predict_rlm.backends.jspi.backend as rlm_interpreter - - # Shrink the request-read budget so the test doesn't wait 30s for - # the default fix to fire. Reusing the module-level knob the fix - # introduces keeps this test honest against future tuning. - if hasattr(rlm_interpreter, "DENO_REQUEST_TIMEOUT_SEC"): - original = rlm_interpreter.DENO_REQUEST_TIMEOUT_SEC - rlm_interpreter.DENO_REQUEST_TIMEOUT_SEC = 0.3 - else: - original = None # RED state — attribute doesn't exist yet - - try: - interp, write_fd = _make_interp_with_partial_pipe() - # Silence writes to stdin so _write_stdin doesn't raise. - class _QuietStdin: - def write(self, _data): pass - def flush(self): pass - def close(self): pass - interp.deno_process = types.SimpleNamespace( - stdin=_QuietStdin(), - stdout=types.SimpleNamespace(fileno=lambda: interp._stdout_fd), - stderr=None, - poll=lambda: None, - ) - interp._request_id = 0 - interp._use_jspi = False - interp._stdin_fd = -1 # force the stdin.write() fallback path - interp._loop = None - - result = {"returned": False, "exc": None, "elapsed": 0.0} - - def _call(): - t0 = time.monotonic() - try: - interp._send_request("health_check", {}, context="test") - except BaseException as e: - result["exc"] = e - result["elapsed"] = time.monotonic() - t0 - result["returned"] = True - - t = threading.Thread(target=_call, daemon=True) - t.start() - t.join(timeout=3.0) - - assert not t.is_alive(), ( - "JspiBackend._send_request hung on a silent deno stdout — " - "_send_request is still passing timeout=None to _read_with_timeout" - ) - assert result["exc"] is not None, ( - "_send_request returned without raising — the silent-stdout " - "case should produce a CodeInterpreterError, not a silent success" - ) - assert result["elapsed"] < 2.0, ( - f"_send_request took {result['elapsed']:.2f}s to fail — " - "timeout enforcement is too loose" - ) - finally: - if original is not None: - rlm_interpreter.DENO_REQUEST_TIMEOUT_SEC = original - try: - os.close(write_fd) - except OSError: - pass - try: - os.close(interp._stdout_fd) - except OSError: - pass - - -def test_read_line_raw_returns_full_line_when_newline_arrives(): - """Healthy-path guardrail: when the writer eventually emits a newline, - ``_read_line_raw`` must return the full line. Guards against the fix - regressing to "always raise/return None prematurely". - """ - interp, write_fd = _make_interp_with_partial_pipe() - try: - # Finish the partial line; writer then closes. - os.write(write_fd, b"-completion\nleftover-bytes") - os.close(write_fd) - write_fd = -1 - - line = interp._read_line_raw() - assert line == "partial-line-no-newline-here-completion" - finally: - if write_fd >= 0: - try: - os.close(write_fd) - except OSError: - pass - try: - os.close(interp._stdout_fd) - except OSError: - pass diff --git a/tests/test_response_id_resync.py b/tests/test_response_id_resync.py index f3be2d41..9a5799a8 100644 --- a/tests/test_response_id_resync.py +++ b/tests/test_response_id_resync.py @@ -1,52 +1,13 @@ -"""RED-GREEN repro for the Response-ID desync bomb. - -Background: - ``JspiBackend._send_request`` increments ``self._request_id``, - writes a JSON-RPC message to deno's stdin, then reads one line from - stdout and asserts the response's ``id`` matches the request. If it - doesn't, the helper raises ``Response ID mismatch``. - - That assumption breaks when a previous ``_send_request`` hit its - ``DENO_REQUEST_TIMEOUT_SEC`` budget and returned ``None``: the deno - process may still deliver its response later, leaving a stale - JSON-RPC frame sitting in the stdout buffer. The NEXT - ``_send_request`` writes a fresh request, reads the stale response - first, sees an older id, and raises. - - The raised error is fed back to the RLM as ``[Error] Response ID - mismatch …``, which the model interprets as a code-format problem. - Since the code isn't broken, the model resubmits identical code - and gets the same mismatch error on the next iteration, burning - through the 50-iteration budget until the task hits - ``task_timeout=600s``. A 2026-04-18 eval run turned 15 tasks into - infinite retry bombs this way — each recorded as a score-0 timeout. - -Fix (option B): treat a non-matching id as a STALE response, discard - it, and read again. Only raise if we exhaust a safety cap on the - number of stale frames or if reading fails. Natural resync; no - fragile buffer-purge logic. - -RED: stale frame preceding the right one → ``Response ID mismatch`` raised -GREEN: stale frame discarded, real response returned cleanly -""" - from __future__ import annotations import asyncio import json -import queue import types import pytest - -pytest.importorskip("websockets") # SBX/supervisor backend requires the [sbx] extra - -pytestmark = pytest.mark.sbx - from dspy.primitives.code_interpreter import CodeInterpreterError # noqa: E402 -import predict_rlm.backends.jspi.backend as rlm_interpreter # noqa: E402 -from predict_rlm.backends import JspiBackend, SbxBackend, SbxConfig # noqa: E402 +from predict_rlm.backends import JspiBackend from predict_rlm.backends.base import STALE_RESPONSE_DISCARD_LIMIT # noqa: E402 @@ -65,10 +26,13 @@ def _build_interp(stdout_lines: list[str]): class _QuietStdin: def __init__(self): self.writes = [] + def write(self, data): self.writes.append(data) + def flush(self): pass + def close(self): pass @@ -90,44 +54,6 @@ def _mock_read(timeout): return interp -def test_stale_response_is_discarded_then_fresh_response_is_returned(monkeypatch): - """When a stale id=5 frame is sitting in the stdout buffer ahead of - the fresh id=7 response, ``_send_request`` must skip the stale and - return the fresh — NOT raise Response ID mismatch. - - Without the fix: first readline returns the stale id=5 frame, the - id check fires, CodeInterpreterError is raised. The model then - sees the error as a code-format bug and burns its iteration budget - retrying identical code. - """ - # Keep test fast — don't let the real DENO_REQUEST_TIMEOUT_SEC=30s - # stretch into a test failure if something unexpected blocks. - monkeypatch.setattr(rlm_interpreter, "DENO_REQUEST_TIMEOUT_SEC", 0.5) - - stale_line = json.dumps({ - "jsonrpc": "2.0", - "id": 5, - "result": {"output": "[Error] timed out on iter 5"}, - }) + "\n" - fresh_line = json.dumps({ - "jsonrpc": "2.0", - "id": 7, - "result": {"output": "ok"}, - }) + "\n" - - interp = _build_interp([stale_line, fresh_line]) - - # Must NOT raise Response ID mismatch. The resync loop discards - # id=5 and returns the id=7 payload. - result = interp._send_request("test_method", {}, context="unit-test") - - assert result is not None, "expected a result dict; got None" - assert result.get("result", {}).get("output") == "ok", ( - f"expected fresh id=7 response (output=ok), got {result!r}. " - "If this is the stale id=5 response the resync is broken." - ) - - def test_multiple_stale_responses_are_discarded(): """If the buffer holds several stale frames (e.g. the process recovered from multiple timeouts in a row), the resync loop must @@ -156,16 +82,6 @@ def test_exhausted_resync_raises_cleanly(): interp._send_request("m", {}, context="t") -def test_matching_response_passes_through_unchanged(): - """Guardrail: when the first readline IS the matching response - (normal case), the resync loop returns it without discarding. - """ - good = json.dumps({"jsonrpc": "2.0", "id": 7, "result": {"output": "hi"}}) + "\n" - interp = _build_interp([good]) - result = interp._send_request("m", {}, context="t") - assert result.get("result", {}).get("output") == "hi" - - def _build_execute_loop_interp(stdout_lines: list[str]): interp = JspiBackend.__new__(JspiBackend) interp._pending_file_ops = {} @@ -230,12 +146,14 @@ async def test_jspi_execute_loop_routes_file_operation_response_before_resync(): @pytest.mark.asyncio async def test_jspi_execute_loop_routes_tool_calls_without_counting_them_stale(): tool_calls = [ - json.dumps({ - "jsonrpc": "2.0", - "method": "tool_call", - "params": {"name": "tool", "args": [], "kwargs": {}}, - "id": f"tool-{idx}", - }) + json.dumps( + { + "jsonrpc": "2.0", + "method": "tool_call", + "params": {"name": "tool", "args": [], "kwargs": {}}, + "id": f"tool-{idx}", + } + ) for idx in range(STALE_RESPONSE_DISCARD_LIMIT + 1) ] fresh = json.dumps({"jsonrpc": "2.0", "id": 7, "result": {"output": "fresh"}}) @@ -257,91 +175,3 @@ async def _wait_all(pending): assert result == "fresh" assert len(called) == STALE_RESPONSE_DISCARD_LIMIT + 1 - - -class _BufferingStdin: - def __init__(self) -> None: - self.data: list[str] = [] - - def write(self, data: str) -> None: - self.data.append(data) - - def flush(self) -> None: - return None - - -def _build_sbx_request_interp(tmp_path, stdout_lines: list[str]) -> SbxBackend: - interp = SbxBackend( - config=SbxConfig(name="resync-test", exec_timeout=1), - preinstall_packages=False, - _runner_command=["unused"], - _staging_root=tmp_path / "staging", - ) - interp._ensure_process_for_method = lambda method: None # type: ignore[method-assign] - interp._proc = types.SimpleNamespace( - stdin=_BufferingStdin(), - stdout=types.SimpleNamespace(), - stderr=None, - poll=lambda: None, - ) - interp._stdout_lines = queue.Queue() - for line in stdout_lines: - interp._stdout_lines.put(line) - return interp - - -def _close_sbx_request_interp(interp: SbxBackend) -> None: - interp._proc = None - - -def test_sbx_send_request_discards_stale_top_level_response(tmp_path): - stale = json.dumps({"jsonrpc": "2.0", "id": 5, "result": {"output": "stale"}}) - fresh = json.dumps({"jsonrpc": "2.0", "id": 1, "result": {"output": "fresh"}}) - interp = _build_sbx_request_interp(tmp_path, [stale, fresh]) - - try: - result = interp._send_request("execute", {"code": "print('fresh')"}) - finally: - _close_sbx_request_interp(interp) - - assert result["result"]["output"] == "fresh" - - -def test_sbx_send_request_exhausted_resync_raises_cleanly(tmp_path): - stale = json.dumps({"jsonrpc": "2.0", "id": 5, "result": {"output": "stale"}}) - interp = _build_sbx_request_interp( - tmp_path, - [stale] * (STALE_RESPONSE_DISCARD_LIMIT + 1), - ) - - try: - with pytest.raises(CodeInterpreterError, match="stale|resync"): - interp._send_request("execute", {"code": "print('fresh')"}) - finally: - _close_sbx_request_interp(interp) - - -def test_sbx_send_request_routes_tool_calls_without_counting_them_stale( - tmp_path, monkeypatch -): - tool_calls = [ - json.dumps({ - "jsonrpc": "2.0", - "method": "tool_call", - "params": {"name": "tool", "args": [], "kwargs": {}}, - "id": f"tool-{idx}", - }) - for idx in range(STALE_RESPONSE_DISCARD_LIMIT + 1) - ] - fresh = json.dumps({"jsonrpc": "2.0", "id": 1, "result": {"output": "fresh"}}) - interp = _build_sbx_request_interp(tmp_path, [*tool_calls, fresh]) - submitted: list[dict] = [] - monkeypatch.setattr(interp, "_submit_tool_call", submitted.append) - - try: - result = interp._send_request("execute", {"code": "print('fresh')"}) - finally: - _close_sbx_request_interp(interp) - - assert result["result"]["output"] == "fresh" - assert len(submitted) == STALE_RESPONSE_DISCARD_LIMIT + 1 diff --git a/tests/test_rlm_gepa.py b/tests/test_rlm_gepa.py index cc008900..d9e7498a 100644 --- a/tests/test_rlm_gepa.py +++ b/tests/test_rlm_gepa.py @@ -1,19 +1,16 @@ from __future__ import annotations -import argparse import asyncio import json -import os import pickle import random +from contextlib import nullcontext from pathlib import Path from types import SimpleNamespace -import dspy import pytest -import rlm_gepa.cli as cli_module -from predict_rlm.telemetry import JsonlTelemetrySink, TelemetryContext, classify_failure +from predict_rlm.telemetry import JsonlTelemetrySink, TelemetryContext from predict_rlm.trace import ( IterationStep, LMFinishMetadata, @@ -24,52 +21,26 @@ TokenUsage, ToolCall, ) -from rlm_gepa import ( - AgentSpec, - EvaluationContext, - OptimizeConfig, - RLMGepaProject, - agent_spec_from_rlm, - build_merge_signature, - build_patch_merge_signature, - build_proposer_for_rlm, - build_proposer_signature, - check_optimization, - run_optimization, -) -from rlm_gepa.cli import apply_optimize_args, run_project_cli -from rlm_gepa.proposer.merge import VALID_STATUSES, RlmMergeProposer -from rlm_gepa.proposer.rlm import ( - ImproveInstructionsGeneric, - PatchMergeInstructionsGeneric, - RLMInstructionProposer, - SelectedCapability, -) -from rlm_gepa.proposer.selection import ( - PatchMergePair, - pick_patch_merge_pair, -) -from rlm_gepa.reporting import stats as stats_report +from rlm_gepa import AgentSpec, OptimizeConfig, RLMGepaProject +from rlm_gepa.cli import run_project_cli +from rlm_gepa.proposer.merge import RlmMergeProposer +from rlm_gepa.proposer.rlm import RLMInstructionProposer +from rlm_gepa.proposer.selection import pick_patch_merge_pair from rlm_gepa.reporting.cost import CostRow, aggregate_costs_from_log, append_cost_rows -from rlm_gepa.reporting.plots import load_plot_data, make_lineage, resolve_plot_output_paths +from rlm_gepa.reporting.plots import load_plot_data from rlm_gepa.reporting.stats import ( candidate_rows, cost_rows, - eval_cost_rows, - eval_task_rows, iteration_rows, merge_rows, render_stats, - render_table, ) from rlm_gepa.runtime.acceptance import should_accept_reflective_candidate -from rlm_gepa.runtime.adapter import RLMGepaAdapter, _row_failure_metadata +from rlm_gepa.runtime.adapter import RLMGepaAdapter from rlm_gepa.schema import RLMGepaExampleResult, validate_project from rlm_gepa.service import ( - GroupAwareBatchSampler, _build_minibatch_sampler, _coerce_reflection_lm_text, - _ProgressCandidateSelector, prepare_run_dir, ) @@ -100,20 +71,6 @@ def _spec() -> AgentSpec: ) -def _spec_without_agent_type() -> AgentSpec: - return AgentSpec( - use_cases=["case a", "case b"], - runtime_grounding_examples={ - "tools": ["tool()"], - "env": ["sandbox timeout"], - "spec": ["protocol behavior"], - }, - tool_signatures="tool() -> str", - target_signature="input: str -> output: str", - scoring_description="score is exact match", - ) - - class _Project(RLMGepaProject): project_name = "test-project" components = ("skill_instructions",) @@ -132,20 +89,6 @@ async def evaluate_example(self, candidate, example, context): # pragma: no cov raise NotImplementedError -def test_reflective_candidate_accepts_bounded_dense_loss_with_hard_flip_signal(): - decision = should_accept_reflective_candidate( - before_scores=[0.99, 0.99, 0.99, 0.50], - after_scores=[1.00, 1.00, 1.00, 0.45], - ) - - assert decision.accepted - assert decision.reason == "hard_flip_signal" - assert decision.dense_delta < 0.0 - assert decision.hard_wins == 3 - assert decision.hard_losses == 0 - assert decision.hard_flip_p_value <= 0.40 - - def test_reflective_candidate_rejects_bounded_dense_loss_without_significant_hard_flips(): decision = should_accept_reflective_candidate( before_scores=[0.99, 0.80], @@ -184,120 +127,6 @@ def test_reflective_candidate_accepts_two_sided_hard_flip_signal_under_default_t assert decision.hard_flip_p_value == pytest.approx(0.375) -def test_build_signatures_render_agent_spec(): - spec = _spec() - proposer = build_proposer_signature(spec) - merge = build_merge_signature(spec) - - assert "test agent" in proposer.instructions - assert "{{" not in proposer.instructions - assert "paired_disagreement_traces_file" in merge.input_fields - - -def test_agent_spec_agent_type_is_optional(): - proposer = build_proposer_signature(_spec_without_agent_type()) - merge = build_merge_signature(_spec_without_agent_type()) - - assert "target RLM" in proposer.instructions - assert "target RLM" in merge.instructions - assert "for ." not in proposer.instructions - assert "{{" not in proposer.instructions - - -def test_agent_spec_from_rlm_can_omit_agent_type(): - class DemoSignature(dspy.Signature): - """Answer questions.""" - - question: str = dspy.InputField() - answer: str = dspy.OutputField() - - def lookup(query: str) -> str: - """Look up a fact.""" - return query - - rlm = SimpleNamespace(signature=DemoSignature, tools=[lookup]) - - spec = agent_spec_from_rlm( - rlm, - use_cases=["case a", "case b"], - runtime_grounding_examples={ - "tools": ["lookup(query)"], - "env": ["sandbox timeout"], - "spec": ["question answering"], - }, - scoring_description="score is exact match", - ) - - assert spec.agent_type == "" - assert "DemoSignature" in spec.target_signature - assert "lookup" in spec.tool_signatures - - -def test_proposer_signature_mentions_parallel_predict_analysis_for_concrete_edits(): - proposer = build_proposer_signature(_spec()) - instructions = " ".join(proposer.instructions.split()) - - assert "predict()" in instructions - assert "asyncio.gather" in instructions - assert "concrete" in instructions - assert "new_instructions" in instructions - - -def test_patch_merge_signature_uses_base_and_patch_source_without_ancestor(): - patch = build_patch_merge_signature(_spec()) - - assert "base_parent_id" in patch.input_fields - assert "base_parent_instructions" in patch.input_fields - assert "patch_source_parent_id" in patch.input_fields - assert "patch_source_parent_instructions" in patch.input_fields - assert "paired_disagreement_traces_file" in patch.input_fields - assert "common_ancestor_instructions" not in patch.input_fields - assert "common ancestor" not in patch.instructions.lower() - - -def test_patch_merge_signature_exposes_selected_capability_contract(): - patch = build_patch_merge_signature(_spec()) - - assert "patch_summary" in patch.output_fields - assert "selected_capability" in patch.output_fields - assert "patch_audit" in patch.output_fields - assert "new_instructions" in patch.output_fields - assert set(SelectedCapability.model_fields) == { - "decision", - "summary", - "evidence_task_ids", - "trigger", - "non_application_boundary", - } - - -def test_merge_signature_is_evidence_backed_patch_contract(): - merge = build_merge_signature(_spec()) - - assert set(merge.input_fields) == set(build_patch_merge_signature(_spec()).input_fields) - assert "base_parent_id" in merge.input_fields - assert "paired_disagreement_traces_file" in merge.input_fields - assert "common_ancestor_instructions" not in merge.input_fields - assert "synthesize" not in merge.instructions.lower() - - -def test_validate_project_accepts_minimal_project(): - result = validate_project(_Project()) - assert result.seed_candidate == {"skill_instructions": "seed rules"} - assert list(result.trainset) == ["train"] - assert list(result.valset) == ["val"] - - -def test_minibatch_sampler_uses_flat_epoch_shuffle_when_project_has_no_group_ids(): - from gepa.core.data_loader import ensure_loader - from gepa.strategies.batch_sampler import EpochShuffledBatchSampler - - loader = ensure_loader(["a", "b", "c", "d"]) - sampler = _build_minibatch_sampler(_Project(), loader, minibatch_size=2, rng=random.Random(7)) - - assert isinstance(sampler, EpochShuffledBatchSampler) - - def test_group_aware_batch_sampler_keeps_groups_intact(): from gepa.core.data_loader import ensure_loader @@ -318,7 +147,6 @@ def minibatch_group_id(self, example) -> str | None: rng=random.Random(7), ) - assert isinstance(sampler, GroupAwareBatchSampler) batch_ids = sampler.next_minibatch_ids(loader, SimpleNamespace(i=0)) batch_examples = loader.fetch(batch_ids) group_counts: dict[str, int] = {} @@ -415,84 +243,6 @@ def test_pick_patch_merge_pair_dedups_sorted_pair_across_ancestors(): assert pair is None -def test_pick_patch_merge_pair_weighted_sampling_is_not_pure_argmax(): - class FakeRng: - def __init__(self): - self.population = [] - self.weights = [] - - def choices(self, population, weights, k): - assert k == 1 - self.population = list(population) - self.weights = list(weights) - return [self.population[1]] - - parents = [[None], [0], [0], [0]] - candidates = [ - {"skill_instructions": "seed"}, - {"skill_instructions": "parent one"}, - {"skill_instructions": "parent two"}, - {"skill_instructions": "parent three"}, - ] - scores_1 = {"t0": 1.0, "t1": 1.0, "t2": 0.0, "t3": 0.0} - scores_2 = {"t0": 0.0, "t1": 0.0, "t2": 1.0, "t3": 1.0} - scores_3 = {"t0": 0.3, "t1": 0.3, "t2": 0.9, "t3": 0.9} - rng = FakeRng() - - pair = pick_patch_merge_pair( - merge_candidates=[1, 2, 3], - program_candidates=candidates, - parent_program_for_candidate=parents, - prog_candidate_val_subscores=[{}, scores_1, scores_2, scores_3], - tracked_scores=[0.0, 0.4, 0.4, 0.4], - merges_performed=[], - rng=rng, - component_name="skill_instructions", - min_each=2, - ) - - assert pair == rng.population[1] - assert len(rng.population) > 1 - weights_by_pair = { - (item.parent_a_id, item.parent_b_id): weight - for item, weight in zip(rng.population, rng.weights, strict=True) - } - assert weights_by_pair[(1, 2)] > weights_by_pair[(1, 3)] - assert (pair.parent_a_id, pair.parent_b_id) != (1, 2) - - -def test_pick_patch_merge_pair_ties_pair_weights_deterministically(): - class NoChoiceRng: - def choices(self, *_args, **_kwargs): - raise AssertionError("equal-weight selector should not call rng.choices") - - parents = [[None], [0], [0], [0]] - candidates = [ - {"skill_instructions": "seed"}, - {"skill_instructions": "parent one"}, - {"skill_instructions": "parent two"}, - {"skill_instructions": "parent three"}, - ] - scores_1 = {"t0": 1.0, "t1": 1.0, "t2": 0.0, "t3": 0.0} - scores_2 = {"t0": 0.0, "t1": 0.0, "t2": 1.0, "t3": 1.0} - scores_3 = dict(scores_2) - - pair = pick_patch_merge_pair( - merge_candidates=[3, 2, 1], - program_candidates=candidates, - parent_program_for_candidate=parents, - prog_candidate_val_subscores=[{}, scores_1, scores_2, scores_3], - tracked_scores=[0.0, 0.4, 0.4, 0.4], - merges_performed=[], - rng=NoChoiceRng(), - component_name="skill_instructions", - min_each=2, - ) - - assert pair is not None - assert (pair.parent_a_id, pair.parent_b_id) == (1, 2) - - def test_pick_patch_merge_pair_chooses_higher_tracked_parent_as_base(): parents = [[None], [0], [0]] candidates = [ @@ -520,125 +270,6 @@ def test_pick_patch_merge_pair_chooses_higher_tracked_parent_as_base(): assert pair.patch_source_parent_id == 2 -def test_pick_patch_merge_pair_ties_base_parent_deterministically(): - parents = [[None], [0], [0]] - candidates = [ - {"skill_instructions": "seed"}, - {"skill_instructions": "lower id"}, - {"skill_instructions": "higher id"}, - ] - scores_a = {f"t{i}": 1.0 if i % 2 else 0.0 for i in range(6)} - scores_b = {f"t{i}": 0.0 if i % 2 else 1.0 for i in range(6)} - - pair = pick_patch_merge_pair( - merge_candidates=[2, 1], - program_candidates=candidates, - parent_program_for_candidate=parents, - prog_candidate_val_subscores=[{}, scores_a, scores_b], - tracked_scores=[0.0, 0.5, 0.5], - merges_performed=[], - rng=random.Random(0), - component_name="skill_instructions", - min_each=2, - ) - - assert pair is not None - assert pair.base_parent_id == 1 - assert pair.patch_source_parent_id == 2 - - -def test_rlm_merge_proposer_uses_patch_selector_by_default(tmp_path: Path, monkeypatch): - from gepa.core.data_loader import ensure_loader - - import rlm_gepa.proposer.merge as merge_module - - calls = {"patch": 0} - - def fake_find_dominators(*_args): - return [1, 2] - - def fake_patch_selector(**_kwargs): - calls["patch"] += 1 - return None - - def evaluator(_inputs, _candidate): - return [], [], None - - monkeypatch.setattr(merge_module, "find_dominator_programs", fake_find_dominators) - monkeypatch.setattr(merge_module, "pick_patch_merge_pair", fake_patch_selector) - - state = SimpleNamespace( - i=0, - full_program_trace=[{}], - program_candidates=[ - {"skill_instructions": "ancestor"}, - {"skill_instructions": "parent a"}, - {"skill_instructions": "parent b"}, - ], - parent_program_for_candidate=[[None], [0], [0]], - prog_candidate_val_subscores=[{}, {"v1": 1.0}, {"v1": 0.0}], - program_full_scores_val_set=[0.0, 0.5, 0.4], - per_program_tracked_scores=[0.0, 0.5, 0.4], - total_num_evals=0, - ) - state.get_pareto_front_mapping = lambda: {} - proposer = RlmMergeProposer( - logger=_Logger(), - valset=ensure_loader(["val"]), - evaluator=evaluator, - adapter=SimpleNamespace(), - trainset=ensure_loader(["train"]), - use_merge=True, - max_merge_invocations=1, - max_rlm_merge_attempts=5, - min_each=1, - merge_minibatch_size=1, - rlm_merge_state_path=tmp_path / "state.json", - rng=random.Random(0), - ) - proposer.last_iter_found_new_program = True - proposer.merges_due = 1 - - assert proposer.propose(state) is None - assert calls == {"patch": 1} - - -def _make_merge_helper(tmp_path: Path) -> RlmMergeProposer: - proposer = RlmMergeProposer.__new__(RlmMergeProposer) - proposer.rlm_merge_state_path = tmp_path / "rlm_merge_state.json" - proposer.rlm_merge_attempts_used = 0 - proposer.merges_performed = ([], []) - return proposer - - -def test_rlm_merge_status_helpers_validate_status_and_namespace(tmp_path: Path): - proposer = _make_merge_helper(tmp_path) - state = SimpleNamespace(full_program_trace=[{"i": 0}]) - - proposer._record_merge_status( - state, - "accepted", - attempt_idx=0, - rlm_merge_candidate_pair=(1, 2), - ) - - assert state.full_program_trace[-1]["rlm_merge_status"] == "accepted" - assert state.full_program_trace[-1]["rlm_merge_candidate_pair"] == (1, 2) - assert VALID_STATUSES == frozenset( - { - "attempt_cap_exhausted", - "pair_skipped", - "preflight_failed", - "subsample_rejected", - "accepted", - "error", - } - ) - - with pytest.raises(ValueError, match="invalid merge status"): - proposer._record_merge_status(state, "not_real", attempt_idx=None) - with pytest.raises(ValueError, match="must start with 'rlm_merge_'"): - proposer._record_merge_status(state, "accepted", attempt_idx=0, bad_field="value") class _FirstKRng(random.Random): def sample(self, population, k): return list(population)[:k] @@ -653,19 +284,14 @@ def __init__(self, tmp_path: Path, base_scores: list[float], source_scores: list self.base_scores = base_scores self.source_scores = source_scores self.evaluate_calls = 0 - self.progress_labels: list[str] = [] - - def progress_label(self, label): - self.progress_labels.append(label) + self.patch_calls = 0 - class NoopContext: - def __enter__(self): - return None + def progress_label(self, _label): + return nullcontext() - def __exit__(self, *_args): - return False - - return NoopContext() + def _rlm_propose_patch_merge_texts(self, **_kwargs): + self.patch_calls += 1 + return "patched instructions", {"patch_summary": "imported one clause"} def evaluate(self, batch, _candidate, *, capture_traces, kind): scores = self.base_scores if self.evaluate_calls == 0 else self.source_scores @@ -710,9 +336,15 @@ def _patch_evidence_state(): i=0, full_program_trace=[{}], program_candidates=[ - {"skill_instructions": "ancestor"}, - {"skill_instructions": "base"}, - {"skill_instructions": "source"}, + {"skill_instructions": "ancestor", "other": "ancestor kept"}, + {"skill_instructions": "base", "other": "base kept"}, + {"skill_instructions": "source", "other": "source ignored"}, + ], + parent_program_for_candidate=[[None], [0], [0]], + prog_candidate_val_subscores=[ + {}, + {"v1": 1.0, "v2": 1.0, "v3": 0.0, "v4": 0.0}, + {"v1": 0.0, "v2": 0.0, "v3": 1.0, "v4": 1.0}, ], total_num_evals=0, ) @@ -748,48 +380,6 @@ def evaluator(_inputs, _candidate): return proposer -def test_patch_evidence_oversamples_two_minibatches_before_selecting_records(tmp_path: Path): - proposer = _make_patch_evidence_proposer( - tmp_path, - base_scores=[1.0, 0.9, 0.8, 0.7, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6], - source_scores=[0.1, 0.2, 0.3, 0.4, 1.0, 0.9, 0.8, 0.7, 0.6, 0.5], - merge_minibatch_size=4, - ) - - evidence = proposer._build_patch_disagreement_evidence( - state=_patch_evidence_state(), - iteration=1, - attempt_idx=0, - base_parent_id=1, - patch_source_parent_id=2, - ) - - assert len(evidence.sampled_train_ids) == 8 - assert len(evidence.records) == 4 - - -def test_patch_evidence_progress_labels_use_zero_indexed_iteration(tmp_path: Path): - proposer = _make_patch_evidence_proposer( - tmp_path, - base_scores=[1.0, 0.0], - source_scores=[0.0, 1.0], - merge_minibatch_size=2, - ) - - proposer._build_patch_disagreement_evidence( - state=_patch_evidence_state(), - iteration=4, - attempt_idx=0, - base_parent_id=1, - patch_source_parent_id=2, - ) - - assert proposer.adapter.progress_labels == [ - "Iteration 4 Patch Base Parent #1 Trace", - "Iteration 4 Patch Source Parent #2 Trace", - ] - - def test_patch_evidence_prefers_larger_disagreements_and_caps_records(tmp_path: Path): proposer = _make_patch_evidence_proposer( tmp_path, @@ -840,12 +430,41 @@ def test_patch_evidence_balances_base_and_patch_source_win_directions(tmp_path: assert winners.count("patch_source") == 2 -def test_patch_evidence_tops_up_with_both_success_records(tmp_path: Path): +@pytest.mark.parametrize( + ("base_scores", "source_scores", "min_each"), + [ + ([1.0, 1.0], [0.0, 0.0], 1), + ([1.0, 1.0, 0.0, 0.0], [0.0, 0.0, 1.0, 1.0], 2), + ], + ids=["missing-source-wins", "cap-drops-required-evidence"], +) +def test_patch_merge_rejects_unbalanced_selected_evidence( + tmp_path: Path, base_scores, source_scores, min_each +): + proposer = _make_patch_evidence_proposer( + tmp_path, + base_scores=base_scores, + source_scores=source_scores, + merge_minibatch_size=2, + min_each=min_each, + ) + state = _patch_evidence_state() + + proposal = proposer._propose_patch_merge( + state=state, iteration=0, merge_candidates=[1, 2], tracked_scores=[0.0, 0.5, 0.5] + ) + + assert proposal is None + assert proposer.adapter.patch_calls == 0 + assert state.full_program_trace[-1]["rlm_merge_status"] == "preflight_failed" + + +def test_patch_disagreement_trace_jsonl_contains_patch_schema(tmp_path: Path): proposer = _make_patch_evidence_proposer( tmp_path, - base_scores=[1.0, 0.0, 1.0, 1.0, 1.0, 1.0], - source_scores=[0.0, 1.0, 1.0, 1.0, 1.0, 1.0], - merge_minibatch_size=6, + base_scores=[1.0, 0.0, 1.0], + source_scores=[0.0, 1.0, 1.0], + merge_minibatch_size=3, ) evidence = proposer._build_patch_disagreement_evidence( @@ -855,319 +474,61 @@ def test_patch_evidence_tops_up_with_both_success_records(tmp_path: Path): base_parent_id=1, patch_source_parent_id=2, ) + records = [ + json.loads(line) for line in Path(evidence.paired_trace_path).read_text().splitlines() + ] - winners = [record["winner"] for record in evidence.records] - assert winners.count("base") == 1 - assert winners.count("patch_source") == 1 - assert winners.count("both_success") == 4 - assert len(evidence.records) == 6 + assert {record["task_id"]: record["evidence_role"] for record in records} == { + "train_0": "base_win", + "train_1": "patch_source_win", + "train_2": "both_success_guardrail", + } + assert records[0]["schema_version"] == 1 + assert records[0]["winner"] in {"base", "patch_source"} + assert records[0]["evidence_role"] in {"base_win", "patch_source_win"} + assert records[0]["abs_delta"] == pytest.approx(1.0) + assert records[0]["base_parent_id"] == 1 + assert records[0]["patch_source_parent_id"] == 2 + assert "generated_outputs" not in records[0]["base_parent"] + assert "trace_preview" not in records[0]["base_parent"] + assert records[0]["base_parent"]["traces"] + records_text = json.dumps(records) + assert "candidate_hash" not in records_text + assert "telemetry_ref" not in records_text + assert "events_path" not in records_text + assert "trace_id" not in records_text + assert records[0]["patch_source_parent"]["feedback"] -def test_patch_merge_preflight_fails_without_balanced_disagreement_evidence( - tmp_path: Path, - monkeypatch, +@pytest.mark.parametrize("child_scores", [[1.0, 1.0], [0.0, 1.0]], ids=["improved", "tied"]) +def test_patch_merge_requires_improvement_and_preserves_base_components( + tmp_path: Path, child_scores ): - from gepa.core.data_loader import ensure_loader + proposer = _make_patch_evidence_proposer( + tmp_path, base_scores=[1.0, 0.0], source_scores=[0.0, 1.0], merge_minibatch_size=2 + ) + state = _patch_evidence_state() - import rlm_gepa.proposer.merge as merge_module + def cached_evaluate(candidate, ids, fetch, evaluator): + assert candidate == {"skill_instructions": "patched instructions", "other": "base kept"} + return child_scores, len(child_scores) - class FakeAdapter(_PatchEvidenceAdapter): - def __init__(self): - super().__init__(tmp_path, base_scores=[1.0, 1.0], source_scores=[0.0, 0.0]) - self.patch_calls = 0 + state.cached_evaluate = cached_evaluate + proposal = proposer._propose_patch_merge( + state=state, iteration=0, merge_candidates=[1, 2], tracked_scores=[0.0, 0.5, 0.5] + ) - def _rlm_propose_patch_merge_texts(self, **_kwargs): - self.patch_calls += 1 - return "should not be called", {} - - def evaluator(_inputs, _candidate): - return [], [], None - - adapter = FakeAdapter() - monkeypatch.setattr(merge_module, "find_dominator_programs", lambda *_args: [1, 2]) - monkeypatch.setattr( - merge_module, - "pick_patch_merge_pair", - lambda **_kwargs: PatchMergePair( - parent_a_id=1, - parent_b_id=2, - base_parent_id=1, - patch_source_parent_id=2, - ancestor=0, - common_ancestors=(0,), - oracle_score=1.0, - oracle_gain=0.5, - base_wins=("v1",), - patch_source_wins=("v2",), - weight=0.5, - ), - ) - proposer = RlmMergeProposer( - logger=_Logger(), - valset=ensure_loader(["v1", "v2"]), - evaluator=evaluator, - adapter=adapter, - trainset=ensure_loader(["train_1", "train_2"]), - use_merge=True, - max_merge_invocations=1, - max_rlm_merge_attempts=5, - min_each=1, - merge_minibatch_size=2, - rlm_merge_state_path=tmp_path / "state.json", - rng=_FirstKRng(), - ) - proposer.last_iter_found_new_program = True - proposer.merges_due = 1 - state = SimpleNamespace( - i=0, - full_program_trace=[{}], - program_candidates=[ - {"skill_instructions": "ancestor"}, - {"skill_instructions": "base"}, - {"skill_instructions": "source"}, - ], - parent_program_for_candidate=[[None], [0], [0]], - prog_candidate_val_subscores=[{}, {"v1": 1.0, "v2": 0.0}, {"v1": 0.0, "v2": 1.0}], - program_full_scores_val_set=[0.0, 0.5, 0.5], - per_program_tracked_scores=[0.0, 0.5, 0.5], - total_num_evals=0, - ) - state.get_pareto_front_mapping = lambda: {} - - proposal = proposer.propose(state) - - assert proposal is None - assert adapter.patch_calls == 0 - assert state.full_program_trace[-1]["rlm_merge_status"] == "preflight_failed" - assert state.full_program_trace[-1]["rlm_merge_reject_reason"].startswith( - "insufficient patch disagreement evidence" - ) - assert state.full_program_trace[-1]["rlm_merge_preflight_base_wins"] == 2 - assert state.full_program_trace[-1]["rlm_merge_preflight_patch_source_wins"] == 0 - - -def test_patch_merge_preflight_fails_when_prompt_cap_cannot_carry_min_each( - tmp_path: Path, - monkeypatch, -): - from gepa.core.data_loader import ensure_loader - - import rlm_gepa.proposer.merge as merge_module - - class FakeAdapter(_PatchEvidenceAdapter): - def __init__(self): - super().__init__( - tmp_path, - base_scores=[1.0, 1.0, 0.0, 0.0], - source_scores=[0.0, 0.0, 1.0, 1.0], - ) - self.patch_calls = 0 - - def _rlm_propose_patch_merge_texts(self, **_kwargs): - self.patch_calls += 1 - return "should not be called", {} - - def evaluator(_inputs, _candidate): - return [], [], None - - adapter = FakeAdapter() - monkeypatch.setattr(merge_module, "find_dominator_programs", lambda *_args: [1, 2]) - monkeypatch.setattr( - merge_module, - "pick_patch_merge_pair", - lambda **_kwargs: PatchMergePair( - parent_a_id=1, - parent_b_id=2, - base_parent_id=1, - patch_source_parent_id=2, - ancestor=0, - common_ancestors=(0,), - oracle_score=1.0, - oracle_gain=0.5, - base_wins=("v1", "v2"), - patch_source_wins=("v3", "v4"), - weight=0.5, - ), - ) - proposer = RlmMergeProposer( - logger=_Logger(), - valset=ensure_loader(["v1", "v2", "v3", "v4"]), - evaluator=evaluator, - adapter=adapter, - trainset=ensure_loader(["train_1", "train_2", "train_3", "train_4"]), - use_merge=True, - max_merge_invocations=1, - max_rlm_merge_attempts=5, - min_each=2, - merge_minibatch_size=2, - rlm_merge_state_path=tmp_path / "state.json", - rng=_FirstKRng(), - ) - proposer.last_iter_found_new_program = True - proposer.merges_due = 1 - state = SimpleNamespace( - i=0, - full_program_trace=[{}], - program_candidates=[ - {"skill_instructions": "ancestor"}, - {"skill_instructions": "base"}, - {"skill_instructions": "source"}, - ], - parent_program_for_candidate=[[None], [0], [0]], - prog_candidate_val_subscores=[ - {}, - {"v1": 1.0, "v2": 1.0, "v3": 0.0, "v4": 0.0}, - {"v1": 0.0, "v2": 0.0, "v3": 1.0, "v4": 1.0}, - ], - program_full_scores_val_set=[0.0, 0.5, 0.5], - per_program_tracked_scores=[0.0, 0.5, 0.5], - total_num_evals=0, - ) - state.get_pareto_front_mapping = lambda: {} - - proposal = proposer.propose(state) - - assert proposal is None - assert adapter.patch_calls == 0 - assert state.full_program_trace[-1]["rlm_merge_status"] == "preflight_failed" - assert state.full_program_trace[-1]["rlm_merge_preflight_base_wins"] == 2 - assert state.full_program_trace[-1]["rlm_merge_preflight_patch_source_wins"] == 2 - assert state.full_program_trace[-1]["rlm_merge_selected_base_wins"] == 1 - assert state.full_program_trace[-1]["rlm_merge_selected_patch_source_wins"] == 1 - - -def test_patch_disagreement_trace_jsonl_contains_patch_schema(tmp_path: Path): - proposer = _make_patch_evidence_proposer( - tmp_path, - base_scores=[1.0, 0.0], - source_scores=[0.0, 1.0], - merge_minibatch_size=2, - ) - - evidence = proposer._build_patch_disagreement_evidence( - state=_patch_evidence_state(), - iteration=1, - attempt_idx=0, - base_parent_id=1, - patch_source_parent_id=2, - ) - records = [json.loads(line) for line in Path(evidence.paired_trace_path).read_text().splitlines()] - - assert len(records) == 2 - assert records[0]["schema_version"] == 1 - assert records[0]["winner"] in {"base", "patch_source"} - assert records[0]["evidence_role"] in {"base_win", "patch_source_win"} - assert records[0]["abs_delta"] == pytest.approx(1.0) - assert records[0]["base_parent_id"] == 1 - assert records[0]["patch_source_parent_id"] == 2 - assert "generated_outputs" not in records[0]["base_parent"] - assert "trace_preview" not in records[0]["base_parent"] - assert records[0]["base_parent"]["traces"] - records_text = json.dumps(records) - assert "candidate_hash" not in records_text - assert "telemetry_ref" not in records_text - assert "events_path" not in records_text - assert "trace_id" not in records_text - assert records[0]["patch_source_parent"]["feedback"] - - -def test_patch_mode_proposer_wires_base_source_fields_without_ancestor( - tmp_path: Path, - monkeypatch, -): - from gepa.core.data_loader import ensure_loader - - import rlm_gepa.proposer.merge as merge_module - - class FakeAdapter(_PatchEvidenceAdapter): - def __init__(self): - super().__init__(tmp_path, base_scores=[1.0, 0.0], source_scores=[0.0, 1.0]) - self.patch_kwargs = None - - def _reserve_merge_proposer_call_idx(self): - return 9 - - def _rlm_propose_patch_merge_texts(self, **kwargs): - self.patch_kwargs = kwargs - return "patched instructions", {"patch_summary": "imported one clause"} - - def queue_valset_progress_label(self, _label): - pass - - def evaluator(_inputs, _candidate): - return [], [], None - - adapter = FakeAdapter() - monkeypatch.setattr(merge_module, "find_dominator_programs", lambda *_args: [1, 2]) - monkeypatch.setattr( - merge_module, - "pick_patch_merge_pair", - lambda **_kwargs: PatchMergePair( - parent_a_id=1, - parent_b_id=2, - base_parent_id=1, - patch_source_parent_id=2, - ancestor=0, - common_ancestors=(0,), - oracle_score=1.0, - oracle_gain=0.5, - base_wins=("v1",), - patch_source_wins=("v2",), - weight=0.5, - ), - ) - proposer = RlmMergeProposer( - logger=_Logger(), - valset=ensure_loader(["v1", "v2"]), - evaluator=evaluator, - adapter=adapter, - trainset=ensure_loader(["train_1", "train_2"]), - use_merge=True, - max_merge_invocations=1, - max_rlm_merge_attempts=5, - min_each=1, - merge_minibatch_size=2, - rlm_merge_state_path=tmp_path / "state.json", - rng=_FirstKRng(), - ) - proposer.last_iter_found_new_program = True - proposer.merges_due = 1 - captured_child = {} - state = SimpleNamespace( - i=0, - full_program_trace=[{}], - program_candidates=[ - {"skill_instructions": "ancestor", "other": "ancestor kept"}, - {"skill_instructions": "base", "other": "base kept"}, - {"skill_instructions": "source", "other": "source ignored"}, - ], - parent_program_for_candidate=[[None], [0], [0]], - prog_candidate_val_subscores=[{}, {"v1": 1.0, "v2": 0.0}, {"v1": 0.0, "v2": 1.0}], - program_full_scores_val_set=[0.0, 0.5, 0.5], - per_program_tracked_scores=[0.0, 0.5, 0.5], - total_num_evals=0, - ) - state.get_pareto_front_mapping = lambda: {} - - def cached_evaluate(candidate, ids, fetch, _evaluator): - captured_child.update(candidate) - assert list(ids) == [0, 1] - assert fetch(list(ids)) == ["train_1", "train_2"] - return [0.0, 0.0], 2 - - state.cached_evaluate = cached_evaluate - - proposal = proposer.propose(state) - - assert proposal is None - assert adapter.patch_kwargs is not None - assert adapter.patch_kwargs["base_parent_id"] == 1 - assert adapter.patch_kwargs["base_parent_instructions"] == "base" - assert adapter.patch_kwargs["patch_source_parent_id"] == 2 - assert adapter.patch_kwargs["patch_source_parent_instructions"] == "source" - assert "paired_disagreement_traces_file" in adapter.patch_kwargs - assert "common_ancestor_instructions" not in adapter.patch_kwargs - assert captured_child == {"skill_instructions": "patched instructions", "other": "base kept"} + if sum(child_scores) > 1: + assert proposal.candidate == { + "skill_instructions": "patched instructions", + "other": "base kept", + } + assert proposal.parent_program_ids == [1, 2] + assert state.full_program_trace[-1]["rlm_merge_status"] == "accepted" + else: + assert proposal is None + assert state.full_program_trace[-1]["rlm_merge_status"] == "subsample_rejected" + assert state.program_candidates[1] == {"skill_instructions": "base", "other": "base kept"} def test_reflection_lm_text_normalization_accepts_common_payloads(): @@ -1192,67 +553,13 @@ def test_reflection_lm_text_normalization_accepts_common_payloads(): _coerce_reflection_lm_text({"usage": {"input_tokens": 10}}) -def test_progress_candidate_selector_uses_zero_indexed_iteration(): - class Selector: - def select_candidate_idx(self, _state): - return 4 - - class Adapter: - def __init__(self): - self.context = None - - def set_reflective_progress_context(self, **kwargs): - self.context = kwargs - - adapter = Adapter() - selector = _ProgressCandidateSelector(Selector(), adapter) - - assert selector.select_candidate_idx(SimpleNamespace(i=5, program_candidates=[{}, {}, {}])) == 4 - assert adapter.context == {"iteration": 5, "parent_idx": 4, "child_idx": 3} - - -def _make_merge_proposer(tmp_path: Path, state_payload: dict | None = None) -> RlmMergeProposer: - from gepa.core.data_loader import ensure_loader - - state_path = tmp_path / "rlm_merge_state.json" - if state_payload is not None: - state_path.write_text(json.dumps(state_payload)) - - def evaluator(_inputs, _candidate): - return [], [], None - - return RlmMergeProposer( - logger=_Logger(), - valset=ensure_loader(["val"]), - evaluator=evaluator, - adapter=SimpleNamespace(), - trainset=ensure_loader(["train"]), - use_merge=True, - max_merge_invocations=1, - max_rlm_merge_attempts=5, - min_each=1, - merge_minibatch_size=1, - rlm_merge_state_path=state_path, - rng=random.Random(0), - ) - - -def test_rlm_merge_proposer_loads_sidecar_state(tmp_path: Path): - proposer = _make_merge_proposer( +def test_merge_sidecar_preserves_attempts_and_pairs_on_resume(tmp_path: Path): + proposer = _make_patch_evidence_proposer( tmp_path, - { - "schema_version": 1, - "rlm_merge_attempts_used": 2, - "merges_performed": [[1, 2, 0], [3, 4, 1]], - }, + base_scores=[1.0], + source_scores=[0.0], + merge_minibatch_size=1, ) - - assert proposer.rlm_merge_attempts_used == 2 - assert proposer.merges_performed[0] == [(1, 2, 0), (3, 4, 1)] - - -def test_rlm_merge_proposer_flushes_sidecar_on_propose_exit(tmp_path: Path): - proposer = _make_merge_proposer(tmp_path) proposer.use_merge = False proposer.rlm_merge_attempts_used = 1 proposer.merges_performed[0].append((1, 2, 0)) @@ -1260,337 +567,14 @@ def test_rlm_merge_proposer_flushes_sidecar_on_propose_exit(tmp_path: Path): assert proposer.propose(state) is None - data = json.loads((tmp_path / "rlm_merge_state.json").read_text()) - assert data["schema_version"] == 1 - assert data["rlm_merge_attempts_used"] == 1 - assert data["merges_performed"] == [[1, 2, 0]] - assert data["flushed_at"] - assert state.full_program_trace[-1]["invoked_merge"] is True - - -def test_plot_output_paths_default_to_run_dir(tmp_path: Path): - score_path, lineage_path = resolve_plot_output_paths(tmp_path) - - assert score_path == tmp_path / "plots" / "score_vs_rollouts.png" - assert lineage_path == tmp_path / "plots" / "candidate_lineage.png" - - -def test_plot_output_paths_accept_directory_or_prefix(tmp_path: Path): - score_path, lineage_path = resolve_plot_output_paths(tmp_path, tmp_path / "plots") - - assert score_path == tmp_path / "plots" / "score_vs_rollouts.png" - assert lineage_path == tmp_path / "plots" / "candidate_lineage.png" - - score_path, lineage_path = resolve_plot_output_paths(tmp_path, tmp_path / "summary.png") - - assert score_path == tmp_path / "summary_score_vs_rollouts.png" - assert lineage_path == tmp_path / "summary_candidate_lineage.png" - - -class _FakePlotlyFigure: - def __init__(self): - self.traces = [] - self.annotations = [] - self.layout = {} - - def add_trace(self, trace): - self.traces.append(trace) - - def add_annotation(self, **kwargs): - self.annotations.append(kwargs) - - def update_layout(self, **kwargs): - self.layout.update(kwargs) - - -def _fake_plotly_scatter(**kwargs): - return SimpleNamespace(**kwargs) - - -_fake_plotly_go = SimpleNamespace(Figure=_FakePlotlyFigure, Scatter=_fake_plotly_scatter) - - -def test_lineage_draws_all_valid_merge_parent_edges(): - data = { - "n": 4, - "scores": [0.1, 0.2, 0.3, 0.4], - "parents": [[None], [0], [0], [1, 2]], - "best_idx": 3, - "pareto_map": {}, - } - - fig = make_lineage(data, _fake_plotly_go) - edge_trace = fig.traces[0] - node_trace = fig.traces[1] - coord_to_candidate = { - (x, y): int(str(text).split("
", maxsplit=1)[0]) - for x, y, text in zip(node_trace.x, node_trace.y, node_trace.text, strict=True) - } - edges = { - ( - coord_to_candidate[(edge_trace.x[i], edge_trace.y[i])], - coord_to_candidate[(edge_trace.x[i + 1], edge_trace.y[i + 1])], - ) - for i in range(0, len(edge_trace.x), 3) - } - - assert edge_trace.x.count(None) == 4 - assert edges == {(0, 1), (0, 2), (1, 3), (2, 3)} - assert "Parents: 1, 2" in node_trace.hovertext[3] - - -def test_lineage_reflows_merge_candidates_to_reduce_crossings(): - data = { - "n": 8, - "scores": [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8], - "parents": [[None], [0], [0], [1, 2], [0], [4], [4], [1, 4]], - "best_idx": 7, - "pareto_map": {}, - } - - fig = make_lineage(data, _fake_plotly_go) - positions = _lineage_node_positions(fig) - edges = _lineage_edges(fig) - baseline_positions = _primary_tree_lineage_positions(data["parents"]) - baseline_crossings = _edge_crossing_count(edges, baseline_positions) - - assert baseline_crossings > 0 - assert _edge_crossing_count(edges, positions) < baseline_crossings - assert _distance_from_parent_center(3, (1, 2), positions) < _distance_from_parent_center( - 3, (1, 2), baseline_positions - ) - - -def test_best_lineage_annotation_reserves_horizontal_space_from_same_layer_nodes(): - data = { - "n": 8, - "scores": [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8], - "parents": [[None], [0], [0], [1, 2], [0], [4], [4], [1, 4]], - "best_idx": 7, - "pareto_map": {}, - } - - fig = make_lineage(data, _fake_plotly_go) - positions = _lineage_node_positions(fig) - best_annotation = next(annotation for annotation in fig.annotations if "Candidate 7 (best)" in annotation["text"]) - xmin, xmax = _annotation_horizontal_footprint(best_annotation) - same_layer_nodes = { - candidate: x - for candidate, (x, y) in positions.items() - if candidate != 7 and y == positions[7][1] - } - - assert { - candidate: x - for candidate, x in same_layer_nodes.items() - if xmin < x < xmax - } == {} - - -def test_lineage_spaces_nodes_away_from_unowned_straight_edges(): - data = { - "n": 13, - "scores": [0.899, 0.895, 0.934, 0.898, 0.873, 0.945, 0.931, 0.900, 0.925, 0.951, 0.952, 0.910, 0.970], - "parents": [[None], [0], [1], [0], [0], [2], [0], [0], [5], [8, 7], [8], [5], [2]], - "best_idx": 12, - "pareto_map": {}, - } - - fig = make_lineage(data, _fake_plotly_go) - edge_trace = fig.traces[0] - - assert edge_trace.x.count(None) == 13 - assert _lineage_edge_node_intersections(fig) == [] - - -def _annotation_horizontal_footprint(annotation: dict[str, object]) -> tuple[float, float]: - x = float(annotation["x"]) - width = 1.5 - if annotation["xanchor"] == "right": - return x - width, x - return x, x + width - - -def _lineage_node_positions(fig) -> dict[int, tuple[float, float]]: - node_trace = fig.traces[1] - return { - int(str(text).split("
", maxsplit=1)[0]): (float(x), float(y)) - for x, y, text in zip(node_trace.x, node_trace.y, node_trace.text, strict=True) - } - - -def _lineage_edges(fig) -> list[tuple[int, int]]: - edge_trace = fig.traces[0] - positions = _lineage_node_positions(fig) - candidate_by_position = {position: candidate for candidate, position in positions.items()} - return [ - ( - candidate_by_position[(float(edge_trace.x[index]), float(edge_trace.y[index]))], - candidate_by_position[(float(edge_trace.x[index + 1]), float(edge_trace.y[index + 1]))], - ) - for index in range(0, len(edge_trace.x), 3) - ] - - -def _primary_tree_lineage_positions(raw_parents: list[object]) -> dict[int, tuple[float, float]]: - primary_parents = [_first_valid_parent(raw_parent, child) for child, raw_parent in enumerate(raw_parents)] - children: dict[int, list[int]] = {index: [] for index in range(len(raw_parents))} - for child, parent in enumerate(primary_parents): - if parent is not None: - children[parent].append(child) - - depth = [0] * len(raw_parents) - - def compute_depth(node: int, current_depth: int) -> None: - depth[node] = current_depth - for child in children[node]: - compute_depth(child, current_depth + 1) - - for root, parent in enumerate(primary_parents): - if parent is None: - compute_depth(root, 0) - - x_pos: dict[int, float] = {} - next_x = 0 - - def layout(node: int) -> None: - nonlocal next_x - if not children[node]: - x_pos[node] = float(next_x) - next_x += 1 - return - for child in children[node]: - layout(child) - x_pos[node] = sum(x_pos[child] for child in children[node]) / len(children[node]) - - for root, parent in enumerate(primary_parents): - if parent is None: - layout(root) - - return {index: (x_pos[index], float(-depth[index])) for index in range(len(raw_parents))} - - -def _first_valid_parent(raw_parent: object, child: int) -> int | None: - raw_values = raw_parent if isinstance(raw_parent, list | tuple) else [raw_parent] - for value in raw_values: - if isinstance(value, bool) or value is None: - continue - parent = int(value) - if 0 <= parent < child: - return parent - return None - - -def _edge_crossing_count( - edges: list[tuple[int, int]], positions: dict[int, tuple[float, float]] -) -> int: - return sum( - _segments_cross(edge, other_edge, positions) - for index, edge in enumerate(edges) - for other_edge in edges[index + 1 :] - ) - - -def _lineage_edge_node_intersections(fig, node_radius: float = 0.25) -> list[tuple[tuple[float, float], tuple[float, float], int]]: - positions = _lineage_node_positions(fig) - edge_trace = fig.traces[0] - intersections = [] - polyline: list[tuple[float, float]] = [] - for x, y in zip(edge_trace.x, edge_trace.y, strict=True): - if x is None or y is None: - intersections.extend(_polyline_node_intersections(polyline, positions, node_radius)) - polyline = [] - else: - polyline.append((float(x), float(y))) - intersections.extend(_polyline_node_intersections(polyline, positions, node_radius)) - return intersections - - -def _polyline_node_intersections( - polyline: list[tuple[float, float]], - positions: dict[int, tuple[float, float]], - node_radius: float, -) -> list[tuple[tuple[float, float], tuple[float, float], int]]: - if len(polyline) < 2: - return [] - endpoints = {polyline[0], polyline[-1]} - return [ - (start, end, candidate) - for start, end in zip(polyline, polyline[1:]) - for candidate, position in positions.items() - if position not in endpoints and _point_segment_distance(position, start, end) < node_radius - ] - - -def _point_segment_distance( - point: tuple[float, float], start: tuple[float, float], end: tuple[float, float] -) -> float: - px, py = point - sx, sy = start - ex, ey = end - dx = ex - sx - dy = ey - sy - segment_len_sq = dx * dx + dy * dy - if segment_len_sq == 0: - return ((px - sx) ** 2 + (py - sy) ** 2) ** 0.5 - t = max(0.0, min(1.0, ((px - sx) * dx + (py - sy) * dy) / segment_len_sq)) - nearest_x = sx + t * dx - nearest_y = sy + t * dy - return ((px - nearest_x) ** 2 + (py - nearest_y) ** 2) ** 0.5 - - -def _segments_cross( - edge: tuple[int, int], - other_edge: tuple[int, int], - positions: dict[int, tuple[float, float]], -) -> bool: - if set(edge) & set(other_edge): - return False - a, b = positions[edge[0]], positions[edge[1]] - c, d = positions[other_edge[0]], positions[other_edge[1]] - ab_c = _orientation(a, b, c) - ab_d = _orientation(a, b, d) - cd_a = _orientation(c, d, a) - cd_b = _orientation(c, d, b) - return ab_c * ab_d < 0 and cd_a * cd_b < 0 - - -def _orientation( - a: tuple[float, float], b: tuple[float, float], c: tuple[float, float] -) -> float: - return (b[0] - a[0]) * (c[1] - a[1]) - (b[1] - a[1]) * (c[0] - a[0]) - - -def _distance_from_parent_center( - candidate: int, parents: tuple[int, ...], positions: dict[int, tuple[float, float]] -) -> float: - parent_center = sum(positions[parent][0] for parent in parents) / len(parents) - return abs(positions[candidate][0] - parent_center) - - -def test_cost_aggregation_raw_and_logical(tmp_path: Path): - path = tmp_path / "cost_log.jsonl" - row = CostRow( - event_id="event_1", - operation_id="op_1", - attempt_id="attempt_1", - event="minibatch", - role="executor", - model="dummy", - calls=2, - input_tokens=10, - output_tokens=5, - cost_usd=0.1, + resumed = _make_patch_evidence_proposer( + tmp_path, + base_scores=[1.0], + source_scores=[0.0], + merge_minibatch_size=1, ) - append_cost_rows(path, [row, row]) - - raw = aggregate_costs_from_log(path) - logical = aggregate_costs_from_log(path, logical=True) - - assert raw[0].calls == 4 - assert raw[0].cost_usd == pytest.approx(0.2) - assert logical[0].calls == 2 - assert logical[0].cost_usd == pytest.approx(0.1) + assert resumed.rlm_merge_attempts_used == 1 + assert resumed.merges_performed[0] == [(1, 2, 0)] def test_logical_cost_keeps_resumed_operations_with_reused_local_counters(tmp_path: Path): @@ -1625,6 +609,13 @@ def test_logical_cost_keeps_resumed_operations_with_reused_local_counters(tmp_pa assert logical[0].calls == 5 assert logical[0].cost_usd == pytest.approx(0.3) + raw = aggregate_costs_from_log(path) + assert raw[0].calls == 7 + assert raw[0].cost_usd == pytest.approx(0.4) + total = next(row for row in cost_rows(tmp_path) if row["scope"] == "TOTAL") + assert total["total_cost"] == "$0.40" + assert total["repeat_cost"] == "$0.10" + assert total["effective_cost"] == "$0.30" def test_logical_cost_does_not_collapse_legacy_rows_without_operation_ids(tmp_path: Path): @@ -1691,123 +682,34 @@ def test_merge_iteration_rows_use_best_actual_parent_instead_of_oracle(tmp_path: assert merge_stats[0]["hard: best(par) -> merge"] == "0.250 → 0.000 -0.250; 1 → 0" assert merge_stats[0]["flips"] == "+0/-1 -1" assert merge_stats[0]["p"] == "1.00" - assert "score Δ" not in merge_stats[0] -def test_stats_hard_flip_p_values_use_two_sided_exact_for_ties(tmp_path: Path): - parent_scores = [0.0, 0.0, 0.0, 0.0, 1.0, 1.0, 1.0, 1.0] - child_scores = [1.0, 1.0, 1.0, 1.0, 0.0, 0.0, 0.0, 0.0] +def test_merge_rows_use_subsample_scores_for_table_metrics_not_full_val_details(tmp_path: Path): state = { + "prog_candidate_val_subscores": [ + {"a": 1.0, "b": 1.0, "c": 1.0}, + {"a": 1.0, "b": 0.0, "c": 1.0}, + {"a": 1.0, "b": 0.0, "c": 0.0}, + ], "full_program_trace": [ { - "i": 0, - "selected_program_candidate": 0, - "new_program_idx": 1, - "subsample_scores": parent_scores, - "new_subsample_scores": child_scores, - }, - { - "i": 1, + "i": 12, "rlm_merge_candidate_pair": (0, 1), - "id1_subsample_scores": parent_scores, - "id2_subsample_scores": [0.0] * len(parent_scores), - "new_program_subsample_scores": child_scores, - }, + "rlm_merge_ancestor": 0, + "rlm_merge_status": "accepted", + "rlm_merge_base_parent": 0, + "rlm_merge_patch_source_parent": 1, + "rlm_merge_new_program_idx": 2, + "id1_subsample_scores": [0.0, 0.0], + "id2_subsample_scores": [1.0, 0.0], + "new_program_subsample_scores": [1.0, 1.0], + } ], } with (tmp_path / "gepa_state.bin").open("wb") as f: pickle.dump(state, f) - rows = iteration_rows(tmp_path) - merge_stats = merge_rows(tmp_path) - - assert rows[0]["flips"] == "+4/-4 +0" - assert rows[0]["p"] == "1.00" - assert merge_stats[0]["flips"] == "+4/-4 +0" - assert merge_stats[0]["p"] == "1.00" - - -def test_merge_rows_report_accepted_child_without_full_val_detail(tmp_path: Path): - state = { - "prog_candidate_val_subscores": [ - {"a": 0.0, "b": 1.0, "c": 0.0}, - {"a": 1.0, "b": 1.0, "c": 0.0}, - {"a": 0.0, "b": 1.0, "c": 1.0}, - {"a": 1.0, "b": 1.0, "c": 1.0}, - ], - "full_program_trace": [ - { - "i": 8, - "rlm_merge_candidate_pair": (1, 2), - "rlm_merge_ancestor": 0, - "rlm_merge_status": "accepted", - "rlm_merge_base_parent": 1, - "rlm_merge_patch_source_parent": 2, - "new_program_idx": 3, - "id1_subsample_scores": [1.0, 0.0], - "id2_subsample_scores": [0.0, 1.0], - "new_program_subsample_scores": [1.0, 1.0], - } - ], - } - with (tmp_path / "gepa_state.bin").open("wb") as f: - pickle.dump(state, f) - - rows = merge_rows(tmp_path) - - assert [key for key in rows[0] if not key.startswith("_")] == [ - "iter", - "pair@anc", - "soft: best(par) -> merge", - "hard: best(par) -> merge", - "flips", - "p", - "outcome", - ] - assert "pre" not in rows[0] - assert "n" not in rows[0] - assert "val Δ" not in rows[0] - assert "score Δ" not in rows[0] - assert "status" not in rows[0] - assert rows[0]["soft: best(par) -> merge"] == "0.500 → 1.000 +0.500" - assert rows[0]["hard: best(par) -> merge"] == "0.500 → 1.000 +0.500; 1 → 2" - assert rows[0]["flips"] == "+1/-0 +1" - assert rows[0]["p"] == "1.00" - assert rows[0]["outcome"] == "accepted" - assert rows[0]["_muted_prefix"] == { - "soft: best(par) -> merge": "0.500 → 1.000", - "hard: best(par) -> merge": "0.500 → 1.000", - "flips": "+1/-0", - } - assert rows[0]["_detail"] == "→ cand 3" - - -def test_merge_rows_use_subsample_scores_for_table_metrics_not_full_val_details(tmp_path: Path): - state = { - "prog_candidate_val_subscores": [ - {"a": 1.0, "b": 1.0, "c": 1.0}, - {"a": 1.0, "b": 0.0, "c": 1.0}, - {"a": 1.0, "b": 0.0, "c": 0.0}, - ], - "full_program_trace": [ - { - "i": 12, - "rlm_merge_candidate_pair": (0, 1), - "rlm_merge_ancestor": 0, - "rlm_merge_status": "accepted", - "rlm_merge_base_parent": 0, - "rlm_merge_patch_source_parent": 1, - "rlm_merge_new_program_idx": 2, - "id1_subsample_scores": [0.0, 0.0], - "id2_subsample_scores": [1.0, 0.0], - "new_program_subsample_scores": [1.0, 1.0], - } - ], - } - with (tmp_path / "gepa_state.bin").open("wb") as f: - pickle.dump(state, f) - - rows = merge_rows(tmp_path) + rows = merge_rows(tmp_path) assert rows[0]["soft: best(par) -> merge"] == "0.500 → 1.000 +0.500" assert rows[0]["hard: best(par) -> merge"] == "0.500 → 1.000 +0.500; 1 → 2" @@ -1845,90 +747,14 @@ def test_merge_rows_do_not_use_normal_mutation_child_for_rejected_merge_val(tmp_ rows = merge_rows(tmp_path) - assert "val Δ" not in rows[0] - assert "score Δ" not in rows[0] - assert "status" not in rows[0] assert rows[0]["soft: best(par) -> merge"] == "0.500 → 0.500 +0.000" assert rows[0]["hard: best(par) -> merge"] == "0.500 → 0.500 +0.000; 1 → 1" assert rows[0]["flips"] == "+1/-1 +0" assert rows[0]["p"] == "1.00" assert rows[0]["outcome"] == "rejected" - assert rows[0]["_muted_prefix"] == { - "soft: best(par) -> merge": "0.500 → 0.500", - "hard: best(par) -> merge": "0.500 → 0.500", - "flips": "+1/-1", - } assert rows[0]["_detail"] == "not better than best parent" -def test_merge_rows_use_explicit_merge_child_for_accepted_detail(tmp_path: Path): - state = { - "prog_candidate_val_subscores": [ - {"a": 0.0, "b": 1.0, "c": 0.0}, - {"a": 1.0, "b": 1.0, "c": 0.0}, - {"a": 0.0, "b": 1.0, "c": 1.0}, - {"a": 0.0, "b": 0.0, "c": 0.0}, - {"a": 1.0, "b": 1.0, "c": 1.0}, - ], - "full_program_trace": [ - { - "i": 10, - "rlm_merge_candidate_pair": (1, 2), - "rlm_merge_ancestor": 0, - "rlm_merge_status": "accepted", - "rlm_merge_new_program_idx": 4, - "new_program_idx": 3, - "id1_subsample_scores": [1.0, 0.0], - "id2_subsample_scores": [0.0, 1.0], - "new_program_subsample_scores": [1.0, 1.0], - } - ], - } - with (tmp_path / "gepa_state.bin").open("wb") as f: - pickle.dump(state, f) - - rows = merge_rows(tmp_path) - - assert "val Δ" not in rows[0] - assert "score Δ" not in rows[0] - assert "status" not in rows[0] - assert rows[0]["soft: best(par) -> merge"] == "0.500 → 1.000 +0.500" - assert rows[0]["hard: best(par) -> merge"] == "0.500 → 1.000 +0.500; 1 → 2" - assert rows[0]["flips"] == "+1/-0 +1" - assert rows[0]["p"] == "1.00" - assert rows[0]["outcome"] == "accepted" - assert rows[0]["_muted_prefix"] == { - "soft: best(par) -> merge": "0.500 → 1.000", - "hard: best(par) -> merge": "0.500 → 1.000", - "flips": "+1/-0", - } - assert rows[0]["_detail"] == "→ cand 4" - - -def test_iteration_hard_count_transition_aligns_count_width(tmp_path: Path): - parent_scores = [1.0] * 13 + [0.0] * 37 - child_scores = [1.0] * 9 + [0.0] * 41 - state = { - "full_program_trace": [ - { - "i": 8, - "selected_program_candidate": 2, - "subsample_scores": parent_scores, - "new_subsample_scores": child_scores, - } - ] - } - with (tmp_path / "gepa_state.bin").open("wb") as f: - pickle.dump(state, f) - - rows = iteration_rows(tmp_path) - - assert rows[0]["hard: par → child"] == "0.260 → 0.180 -0.080; 13 → 9" - rendered = render_table(rows, width=120) - plain_rendered = stats_report.re.sub(r"\033\[[0-9;]*m", "", rendered) - assert ".260 → .180 -.080; 13 → 9" in plain_rendered - - def test_iteration_rows_include_attempts_without_child_scores(tmp_path: Path): state = { "full_program_trace": [ @@ -1948,275 +774,11 @@ def test_iteration_rows_include_attempts_without_child_scores(tmp_path: Path): rows = iteration_rows(tmp_path) assert [row["iter"] for row in rows] == ["1 [0]", "2 [0]"] - assert rows[0] == { - "iter": "1 [0]", - "soft: par → child": "-", - "hard: par → child": "-", - "flips": "-", - "p": "-", - "outcome": "NO CHILD", - "_highlight": False, - "_muted_prefix": {}, - "_iteration_hard_denominator": 2, - "_terminal_header_suffixes": {"hard: par → child": "/2"}, - } + assert rows[0]["outcome"] == "NO CHILD" + assert rows[0]["soft: par → child"] == "-" assert rows[1]["outcome"] == "REJECTED" -def test_merge_rows_compact_outcomes_for_terminal_width(tmp_path: Path): - state = { - "full_program_trace": [ - {"i": 1, "rlm_merge_status": "pair_skipped", "rlm_merge_candidate_pair": (0, 1)}, - { - "i": 2, - "rlm_merge_status": "no_merge_candidate", - "rlm_merge_candidate_pair": (0, 1), - }, - { - "i": 3, - "rlm_merge_status": "subsample_rejected", - "rlm_merge_candidate_pair": (0, 1), - }, - { - "i": 4, - "rlm_merge_status": "preflight_failed", - "rlm_merge_candidate_pair": (0, 1), - }, - {"i": 5, "rlm_merge_status": "accepted", "rlm_merge_candidate_pair": (0, 1)}, - {"i": 6, "rlm_merge_status": "custom_status", "rlm_merge_candidate_pair": (0, 1)}, - ] - } - with (tmp_path / "gepa_state.bin").open("wb") as f: - pickle.dump(state, f) - - rows = merge_rows(tmp_path) - - assert [row["outcome"] for row in rows] == [ - "skipped", - "skipped", - "rejected", - "rejected", - "accepted", - "custom_status", - ] - - -def test_reporting_tables_from_artifacts(tmp_path: Path): - state = { - "i": 1, - "total_num_evals": 4, - "program_candidates": [{"skill_instructions": "seed"}, {"skill_instructions": "new"}], - "program_full_scores_val_set": [0.5, 0.75], - "prog_candidate_val_subscores": [ - {"a": 0.0, "b": 1.0}, - {"a": 1.0, "b": 1.0}, - ], - "parent_program_for_candidate": [[None], [0]], - "full_program_trace": [ - { - "i": 0, - "selected_program_candidate": 0, - "subsample_scores": {"a": 0.0, "b": 1.0}, - "new_subsample_scores": {"a": 1.0, "b": 1.0}, - "new_program_idx": 1, - }, - { - "i": 1, - "rlm_merge_candidate_pair": (0, 1), - "rlm_merge_ancestor": 0, - "rlm_merge_attempt_idx": 0, - "rlm_merge_status": "subsample_rejected", - "rlm_merge_preflight_a_wins": 3, - "rlm_merge_preflight_b_wins": 2, - "rlm_merge_reject_reason": "not better than best parent", - "id1_subsample_scores": [0.0, 1.0], - "id2_subsample_scores": [1.0, 0.0], - "new_program_subsample_scores": [1.0, 0.0], - }, - ], - } - with (tmp_path / "gepa_state.bin").open("wb") as f: - pickle.dump(state, f) - (tmp_path / "run_metadata.json").write_text( - json.dumps( - { - "resolved_config": { - "executor_reasoning_effort": "low", - "executor_sub_lm_reasoning_effort": "none", - "proposer_reasoning_effort": "medium", - "proposer_sub_lm_reasoning_effort": "medium", - } - } - ) - ) - append_cost_rows( - tmp_path / "cost_log.jsonl", - [ - CostRow( - event_id="e", - operation_id="op", - attempt_id="a", - event="valset", - role="executor", - model="dummy", - calls=1, - input_tokens=10, - output_tokens=2, - cost_usd=0.01, - ), - CostRow( - event_id="e", - operation_id="op", - attempt_id="a", - event="valset", - role="executor", - model="dummy", - calls=1, - input_tokens=10, - output_tokens=2, - cost_usd=0.01, - ), - ], - ) - - rows = iteration_rows(tmp_path) - assert rows[0]["outcome"] == "→ cand 1" - assert rows[0]["soft: par → child"] == "0.500 → 1.000 +0.500" - assert rows[0]["hard: par → child"] == "0.500 → 1.000 +0.500; 1 → 2" - assert rows[0]["flips"] == "+1/-0 +1" - assert rows[0]["p"] == "1.00" - assert rows[0]["iter"] == "0 [0]" - assert rows[0]["_highlight"] is True - iteration_terminal = render_table(rows, width=120) - iteration_plain_lines = [ - stats_report.re.sub(r"\033\[[0-9;]*m", "", line) for line in iteration_terminal.splitlines() - ] - iteration_header_line = next(line for line in iteration_plain_lines if "hard: par" in line) - iteration_hard_header = iteration_header_line.strip("│").split("│")[2] - assert iteration_hard_header.rstrip().endswith("/2") - assert "1 → 2 /2" not in iteration_terminal - assert rows[1]["iter"] == "1 [0, 1]" - merges = merge_rows(tmp_path) - assert { - key: value - for key, value in merges[0].items() - if key not in {"_merge_hard_denominator", "_terminal_header_aliases", "_terminal_header_suffixes"} - } == { - "iter": "1", - "pair@anc": "0+1@0", - "soft: best(par) -> merge": "0.500 → 0.500 +0.000", - "hard: best(par) -> merge": "0.500 → 0.500 +0.000; 1 → 1", - "flips": "+1/-1 +0", - "p": "1.00", - "outcome": "rejected", - "_detail": "not better than best parent", - "_muted_prefix": { - "soft: best(par) -> merge": "0.500 → 0.500", - "hard: best(par) -> merge": "0.500 → 0.500", - "flips": "+1/-1", - }, - } - assert rows[0]["_terminal_header_suffixes"]["hard: par → child"] == "/2" - assert merges[0]["_terminal_header_suffixes"]["hard: best(par) -> merge"] == "/2" - candidates = candidate_rows(tmp_path) - assert candidates[0]["cand [par]"] == "0 [seed]" - assert candidates[0]["soft: par → child"] == "- → 0.500" - assert "soft" not in candidates[0] - assert "mean" not in candidates[0] - assert candidates[0]["hard: par → child"] == "- → 0.500" - assert candidates[0]["flips vs par"] == "-" - assert candidates[0]["Δ-seed"] == "-" - assert candidates[1]["cand [par]"] == "1 [0]" - assert candidates[1]["soft: par → child"] == "0.500 → 1.000 +0.500" - assert "soft" not in candidates[1] - assert "mean" not in candidates[1] - assert candidates[1]["hard: par → child"] == "0.500 → 1.000 +0.500" - assert candidates[1]["flips vs par"] == "+1/-0 +1" - assert candidates[1]["Δ-seed"] == "+0.500" - assert candidates[1]["_muted_prefix"] == { - "soft: par → child": "0.500 → ", - "hard: par → child": "0.500 → ", - "flips vs par": "+1/-0", - } - assert candidates[1]["_muted_suffix"] == { - "soft: par → child": " +0.500", - "hard: par → child": " +0.500", - } - assert candidates[1]["_highlight"] is True - candidate_terminal = render_table(candidates, width=120) - assert "\033[38;5;178m0.500 → \033[0m\033[1;38;5;220m1.000\033[38;5;178m +0.500" in candidate_terminal - costs = cost_rows(tmp_path) - assert costs[0]["scope"] == "executor" - assert costs[0]["model"] == "" - assert costs[0]["calls"] == "" - assert costs[0]["_category"] is True - assert costs[1]["scope"] == " - main" - assert costs[1]["model"] == "dummy-low" - assert costs[1]["total_cost"] == "$0.02" - assert costs[1]["repeat_cost"] == "$0.01" - assert costs[1]["effective_cost"] == "$0.01" - assert costs[2]["scope"] == " - sub" - assert costs[2]["model"] == "-" - assert costs[2]["calls"] == "-" - assert costs[3]["_spacer"] is True - assert costs[-1]["scope"] == "TOTAL" - assert costs[-1]["model"] == "" - assert costs[-1]["calls"] == "2" - assert costs[-1]["total_cost"] == "$0.02" - assert costs[-1]["repeat_cost"] == "$0.01" - assert costs[-1]["effective_cost"] == "$0.01" - rendered = render_stats(tmp_path, output_format="markdown") - assert "iterations:" in rendered - assert "merges:" in rendered - assert "| iter" in rendered - assert "| soft: par → child" in rendered - assert "| hard: par → child" in rendered - assert "| soft: best(par) -> merge" in rendered - assert "| hard: best(par) -> merge" in rendered - assert "| pair@anc" in rendered - assert "rejected" in rendered - assert "merge details:" in rendered - assert "iter 1 0+1@0: not better than best parent" in rendered - assert "| cand [par]" in rendered - candidate_section = rendered.split("candidates:", maxsplit=1)[1].split("costs:", maxsplit=1)[0] - assert "soft: par → child" in candidate_section - assert "hard: par → child" in candidate_section - assert "flips vs par" in candidate_section - assert "| soft |" not in candidate_section - assert "| hard |" not in candidate_section - assert "| flips |" not in candidate_section - assert "| Δ-seed" in rendered - assert "| ----" in rendered - assert "**1 [0]**" in rendered - terminal = render_stats(tmp_path, width=120) - assert "┌" in terminal - assert "\033[3m" in terminal - assert "\033[38;5;248m" in terminal - assert "\033[38;5;178m0.500 → 1.000\033[0m\033[1;38;5;220m +0.500" in terminal - assert "\033[38;5;178m0.500 → 1.000\033[0m\033[1;38;5;220m +0.500; 1 → 2" in terminal - assert "\033[38;5;178m+1/-0\033[0m\033[1;38;5;220m +1" in terminal - assert "\033[38;5;248m+1/-1\033[0m +0" in terminal - assert "\033[38;5;248m0.500 → 0.500\033[0m +0.000" in terminal - assert "\033[38;5;248m0.500 → 0.500\033[0m +0.000; 1 → 1" in terminal - assert "\033[1;38;5;220m" in terminal - assert "**1**" not in terminal - assert "costs:" in terminal - assert "total" in terminal - assert "repeat" in terminal - assert "eff" in terminal - assert "costs (raw spend: all logged LM calls):" not in terminal - assert "costs (deduped spend: stable operation ids only; legacy rows counted raw):" not in terminal - - -def test_render_stats_before_state_checkpoint_does_not_crash(tmp_path: Path) -> None: - (tmp_path / "run_metadata.json").write_text(json.dumps({"project_name": "demo"})) - - rendered = render_stats(tmp_path, table="all", output_format="markdown") - - assert "iter=0" in rendered - assert "candidates=0" in rendered - - def test_candidate_rows_show_flips_against_each_parent(tmp_path: Path): state = { "parent_program_for_candidate": [[None], [0], [0, 1]], @@ -2234,34 +796,10 @@ def test_candidate_rows_show_flips_against_each_parent(tmp_path: Path): assert rows[2]["soft: par → child"] == "0.333 → 0.667 +0.333\n0.333 → 0.667 +0.333" assert rows[2]["hard: par → child"] == "0.333 → 0.667 +0.333\n0.333 → 0.667 +0.333" assert rows[2]["flips vs par"] == "+1/-0 +1\n+1/-0 +1" - assert rows[2]["_muted_prefix"] == { - "soft: par → child": "0.333 → \n0.333 → ", - "hard: par → child": "0.333 → \n0.333 → ", - "flips vs par": "+1/-0\n+1/-0", - } - assert rows[2]["_muted_suffix"] == { - "soft: par → child": " +0.333\n +0.333", - "hard: par → child": " +0.333\n +0.333", - } - - rendered = render_table(rows, width=120) - - assert " -> " not in rendered - assert "\033[38;5;178m.333 → \033[0m\033[1;38;5;220m.667\033[38;5;178m +.333" in rendered - assert "\033[38;5;178m+1/-0\033[0m\033[1;38;5;220m +1" in rendered def test_live_state_best_candidate_overrides_stale_summary(tmp_path: Path): state = { - "full_program_trace": [ - { - "i": 0, - "selected_program_candidate": 0, - "subsample_scores": [0.0, 0.0], - "new_subsample_scores": [1.0, 1.0], - "new_program_idx": 2, - } - ], "program_candidates": [{}, {}, {}], "parent_program_for_candidate": [[None], [0], [1]], "prog_candidate_val_subscores": [ @@ -2283,12 +821,8 @@ def test_live_state_best_candidate_overrides_stale_summary(tmp_path: Path): ) ) - candidates = candidate_rows(tmp_path) - iterations = iteration_rows(tmp_path) plot_data = load_plot_data(tmp_path) - assert [row["_highlight"] for row in candidates] == [False, False, True] - assert iterations[0]["_highlight"] is True assert plot_data["best_idx"] == 2 assert plot_data["scores"] == [0.0, 0.6, 1.0] @@ -2330,316 +864,6 @@ def test_plot_data_repairs_truncated_live_full_scores_from_subscores(tmp_path: P assert plot_data["eval_counts"] == [0, 10, 20, 30, 40, 50, 51, 52] -def test_cost_rows_group_patch_merge_roles(tmp_path: Path): - (tmp_path / "run_metadata.json").write_text( - json.dumps( - { - "resolved_config": { - "proposer_reasoning_effort": "medium", - "proposer_sub_lm_reasoning_effort": "low", - } - } - ) - ) - append_cost_rows( - tmp_path / "cost_log.jsonl", - [ - CostRow( - event_id="e", - operation_id="op", - attempt_id="a", - event="patch_merge", - role="patch_merge_proposer", - model="dummy-patch", - calls=1, - input_tokens=10, - output_tokens=2, - cost_usd=0.01, - ), - CostRow( - event_id="e-sub", - operation_id="op-sub", - attempt_id="a-sub", - event="patch_merge", - role="patch_merge_proposer_sub_lm", - model="dummy-patch-sub", - calls=1, - input_tokens=11, - output_tokens=3, - cost_usd=0.02, - ), - ], - ) - - rows = cost_rows(tmp_path) - - assert any(row.get("scope") == "merge" and row.get("_category") for row in rows) - assert any( - row.get("scope") == " - proposer main" and row.get("model") == "dummy-patch-medium" - for row in rows - ) - assert any( - row.get("scope") == " - proposer sub" and row.get("model") == "dummy-patch-sub-low" - for row in rows - ) - assert not any(row.get("scope") == "patch-merge" for row in rows) - assert not any(row.get("scope") == "other" for row in rows) - - -def test_highlighted_terminal_rows_use_dim_gold_for_muted_prefixes(): - rendered = render_table( - [ - { - "metric": "0.100 → 0.200 +0.100", - "_highlight": True, - "_muted_prefix": {"metric": "0.100 → 0.200"}, - } - ] - ) - - assert "\033[38;5;178m.100 → .200" in rendered - assert "\033[38;5;248m.100 → .200" not in rendered - - -def test_terminal_cost_table_wraps_scope_and_model_to_terminal_width(monkeypatch): - monkeypatch.setattr( - stats_report.shutil, - "get_terminal_size", - lambda fallback=(120, 24): os.terminal_size((82, 24)), - ) - rows = [ - { - "scope": " - patch_merge_proposer_sub_lm", - "model": "openai/gpt-5.4-mini-medium", - "calls": "1,234", - "prompt_tok": "12,345,678", - "completion_tok": "123,456", - "total_cost": "$123.45", - "repeat_cost": "$0.00", - "effective_cost": "$123.45", - } - ] - - rendered = render_table(rows) - plain_lines = [stats_report.re.sub(r"\033\[[0-9;]*m", "", line) for line in rendered.splitlines()] - - assert max(len(line) for line in plain_lines) <= 82 - assert "in_tok" in rendered - assert "out_tok" in rendered - assert "total_cost" not in rendered - assert " - patch" in rendered - assert "│ merge_" in rendered - assert "propos" in rendered - - -def test_terminal_merge_table_wraps_headers_and_status_to_terminal_width(monkeypatch): - def plain_lines(rendered: str) -> list[str]: - return [stats_report.re.sub(r"\033\[[0-9;]*m", "", line) for line in rendered.splitlines()] - - def cell_lines(rendered: str, column_index: int) -> list[str]: - lines = [] - for line in plain_lines(rendered): - if line.startswith("│"): - cells = line.strip("│").split("│") - lines.append(cells[column_index].strip()) - return lines - - def raw_cell_lines(rendered: str, column_index: int) -> list[str]: - lines = [] - for line in plain_lines(rendered): - if line.startswith("│"): - cells = line.strip("│").split("│") - lines.append(cells[column_index]) - return lines - - def body_cell_rows(rendered: str) -> list[list[str]]: - lines = plain_lines(rendered) - body_start = next(index for index, line in enumerate(lines) if line.startswith("├")) + 1 - return [line.strip("│").split("│") for line in lines[body_start:] if line.startswith("│")] - - monkeypatch.setattr( - stats_report.shutil, - "get_terminal_size", - lambda fallback=(120, 24): os.terminal_size((90, 24)), - ) - rows = [ - { - "iter": "12345 [123, 456]", - "pair@anc": "123+456@789", - "soft: best(par) -> merge": "0.123 → 0.987 +0.864", - "hard: best(par) -> merge": "0.111 → 0.999 +0.888; 1 → 9", - "flips": "+8/-0 +8", - "p": "0.01", - "outcome": "rejected", - "_terminal_header_aliases": { - "soft: best(par) -> merge": "soft\nbest(par) -> merge", - "hard: best(par) -> merge": "hard\nbest(par) -> merge", - }, - "_terminal_header_suffixes": { - "hard: best(par) -> merge": "/10", - }, - "_muted_prefix": { - "soft: best(par) -> merge": "0.123 → 0.987", - "hard: best(par) -> merge": "0.111 → 0.999", - "flips": "+8/-0", - }, - } - ] - - rendered = render_table(rows) - - assert "soft: best(par) -> merge" not in rendered - assert "hard: best(par) -> merge" not in rendered - assert "soft" in rendered - assert "best(par) -> merge" in rendered - assert stats_report.re.search(r"best\(par\) -> merge\s+/10", "\n".join(plain_lines(rendered))) - assert "par→merge" not in rendered - assert "1 → 9 /10" not in rendered - assert "rejected" in rendered - - monkeypatch.setattr( - stats_report.shutil, - "get_terminal_size", - lambda fallback=(120, 24): os.terminal_size((92, 24)), - ) - moderate = render_table( - [ - { - "iter": "12", - "pair@anc": "3+4@2", - "soft: best(par) -> merge": "0.123 → 0.987 +0.864", - "hard: best(par) -> merge": "0.111 → 0.999 +0.888; 1 → 9", - "flips": "+8/-0 +8", - "p": "0.01", - "outcome": "accepted", - "_terminal_header_aliases": { - "soft: best(par) -> merge": "soft\nbest(par) -> merge", - "hard: best(par) -> merge": "hard\nbest(par) -> merge", - }, - "_terminal_header_suffixes": { - "hard: best(par) -> merge": "/10", - }, - "_muted_prefix": { - "soft: best(par) -> merge": "0.123 → 0.987", - "hard: best(par) -> merge": "0.111 → 0.999", - "flips": "+8/-0", - }, - } - ] - ) - moderate_lines = plain_lines(moderate) - pair_lines = cell_lines(moderate, 1) - hard_lines = raw_cell_lines(moderate, 3) - moderate_body_cells = body_cell_rows(moderate) - - assert moderate_lines - assert pair_lines[:2] == ["pair", "@anc"] - assert "@" not in pair_lines[:2] - assert "anc" not in pair_lines[:2] - assert pair_lines[2:3] == ["3+4@2"] - assert hard_lines[1].rstrip().endswith("/10") - assert len(moderate_body_cells) == 1 - assert [moderate_body_cells[0][index].strip() for index in (1, 2, 3, 4)] == [ - "3+4@2", - ".123 → .987 +.864", - ".111 → .999 +.888; 1 → 9", - "+8/-0 +8", - ] - - monkeypatch.setattr( - stats_report.shutil, - "get_terminal_size", - lambda fallback=(120, 24): os.terminal_size((90, 24)), - ) - tight = render_table(rows) - tight_lines = plain_lines(tight) - body_cells = body_cell_rows(tight) - - assert tight_lines - assert len(body_cells) == 1 - assert [body_cells[0][index].strip() for index in (1, 2, 3, 4)] == [ - "123+456@789", - ".123 → .987 +.864", - ".111 → .999 +.888; 1 → 9", - "+8/-0 +8", - ] - - -def test_eval_stats_from_eval_artifact(tmp_path: Path): - report = { - "config": {"reasoning_effort": "medium"}, - "total_tasks": 2, - "soft_restriction_avg": 0.75, - "hard_restriction_avg": 0.5, - "tasks_all_passing": 1, - "duration_seconds": 125, - "total_cost_usd": 1.23, - "costs": [ - { - "role": "main", - "model": "dummy-main", - "calls": 3, - "prompt_tokens": 100, - "completion_tokens": 20, - "cost_usd": 1.0, - }, - { - "role": "sub", - "model": "dummy-sub", - "calls": 2, - "prompt_tokens": 10, - "completion_tokens": 5, - "cost_usd": 0.23, - }, - ], - "per_task": [ - { - "task_id": "a", - "soft": 1.0, - "hard": 1, - "cases": [ - {"passed": True, "message": "All 10 cells match"}, - {"passed": True, "message": "All 5 cells match"}, - ], - }, - { - "task_id": "b", - "soft": 0.5, - "hard": 0, - "cases": [ - {"passed": False, "message": "Sheet 'A' range A1:A6: 3/6 cells match"} - ], - }, - ], - } - (tmp_path / "eval.json").write_text(json.dumps(report)) - - tasks = eval_task_rows(tmp_path) - costs = eval_cost_rows(tmp_path) - rendered = render_stats(tmp_path, table="all") - - assert tasks[0] == { - "task": "a", - "soft": "1.000 (15 /15)", - "hard": "1.000 (2 /2)", - "_align": {"soft": "left"}, - } - assert tasks[1] == { - "task": "b", - "soft": "0.500 (3 /6)", - "hard": "0.000 (0 /1)", - "_align": {"soft": "left"}, - } - assert costs[0]["scope"] == "executor" - assert costs[1]["scope"] == " - main" - assert costs[1]["model"] == "dummy-main-medium" - assert costs[2]["scope"] == " - sub" - assert costs[2]["model"] == "dummy-sub-none" - assert "eval: tasks=2, soft=0.750, hard=0.500 (1/2), cost=$1.23, duration=2m 5s" in rendered - assert "tasks:" in rendered - assert "costs:" in rendered - - def test_eval_stats_reports_attempt_outcomes_and_latency_percentiles(tmp_path: Path): report = { "total_tasks": 4, @@ -2661,7 +885,12 @@ def test_eval_stats_reports_attempt_outcomes_and_latency_percentiles(tmp_path: P "example_id": "ok", "status": "completed", "feedback": "passed", - "trace": {"status": "completed", "iterations": 2, "max_iterations": 5, "duration_ms": 1000}, + "trace": { + "status": "completed", + "iterations": 2, + "max_iterations": 5, + "duration_ms": 1000, + }, }, { "example_id": "outer-timeout", @@ -2674,560 +903,60 @@ def test_eval_stats_reports_attempt_outcomes_and_latency_percentiles(tmp_path: P "status": "completed", "feedback": "submitted by fallback", "trace": { - "status": "max_iterations", - "iterations": 5, - "max_iterations": 5, - "duration_ms": 5000, - }, - }, - { - "example_id": "project-timeout", - "status": "error", - "error": "RLM timeout at 300s", - "trace": {"status": "error", "iterations": 1, "max_iterations": 5, "duration_ms": 2000}, - }, - ] - ) - ) - - terminal = render_stats(tmp_path, table="all") - markdown = render_stats(tmp_path, table="all", output_format="markdown") - - expected = ( - "attempts=4, timeouts=2, max_iter_hits=1, " - "latency p50=2.0s p90=5.0s p95=5.0s max=5.0s" - ) - assert expected in terminal - assert expected in markdown - - -def test_optimize_stats_reports_attempt_outcomes_and_latency_percentiles(tmp_path: Path): - with (tmp_path / "gepa_state.bin").open("wb") as f: - pickle.dump( - { - "i": 0, - "program_candidates": [], - "total_num_evals": 2, - "prog_candidate_val_subscores": [], - }, - f, - ) - trace_dir = tmp_path / "task_traces" - trace_dir.mkdir() - (trace_dir / "eval_minibatch_attempts.jsonl").write_text( - "\n".join( - json.dumps(row) - for row in [ - { - "example_id": "a", - "status": "completed", - "trace": {"status": "completed", "iterations": 1, "max_iterations": 3, "duration_ms": 2500}, - }, - { - "example_id": "b", - "status": "completed", - "trace": {"status": "completed", "iterations": 3, "max_iterations": 3, "duration_ms": 7500}, - }, - { - "example_id": "c", - "status": "timeout", - "feedback": "timed out", - "trace": None, - }, - ] - ) - ) - - rendered = render_stats(tmp_path, table="costs") - - assert ( - "attempts=3, timeouts=1, max_iter_hits=1, " - "latency p50=2.5s p90=7.5s p95=7.5s max=7.5s" - ) in rendered - - -def test_render_table_outputs_github_markdown(): - rendered = render_table([{"a": "x|y", "b": "z\nw"}], output_format="markdown") - - assert rendered.splitlines()[0].startswith("| a") - assert rendered.splitlines()[1].startswith("| -") - assert "x\\|y" in rendered - assert "z
w" in rendered - - -def test_render_table_compacts_fractional_decimal_columns(): - rows = [ - {"mean": "0.123", "delta": "+0.045", "mixed": "0.123 → 1.000", "model": "gpt-0.5"}, - {"mean": "0.456", "delta": "-0.012", "mixed": "0.456 → 1.000", "model": "gpt-0.7"}, - ] - - markdown = render_table(rows, output_format="markdown") - terminal = render_table(rows, output_format="terminal", width=120) - - assert ".123" in markdown - assert "+.045" in markdown - assert "-.012" in markdown - assert "0.123 → 1.000" in markdown - assert "gpt-0.5" in markdown - assert ".456" in terminal - assert "gpt-0.7" in terminal - - -def test_project_cli_check_with_dummy_lms(capsys): - config = OptimizeConfig( - executor_lm=_DummyLM(), - executor_sub_lm=_DummyLM(), - proposer_lm=_DummyLM(), - proposer_sub_lm=_DummyLM(), - ) - status = run_project_cli(lambda: _Project(), config, argv=["optimize", "--check"]) - - assert status == 0 - assert "check ok" in capsys.readouterr().out - - -def test_project_cli_accepts_stat_alias(monkeypatch, capsys, tmp_path: Path): - def render_stats(run_dir, table="all", output_format="terminal", width=None): - return f"stats {Path(run_dir).name} {table} {output_format} {width}" - - monkeypatch.setattr(cli_module, "render_stats", render_stats) - - status = run_project_cli(lambda: _Project(), OptimizeConfig(), argv=["stat", str(tmp_path)]) - - assert status == 0 - assert f"stats {tmp_path.name} all terminal None" in capsys.readouterr().out - - -def test_project_cli_passes_stats_width(monkeypatch, capsys, tmp_path: Path): - def render_stats(run_dir, table="all", output_format="terminal", width=None): - return f"stats {Path(run_dir).name} {table} {output_format} {width}" - - monkeypatch.setattr(cli_module, "render_stats", render_stats) - - status = run_project_cli( - lambda: _Project(), - OptimizeConfig(), - argv=["stats", str(tmp_path), "--width", "107"], - ) - - assert status == 0 - assert f"stats {tmp_path.name} all terminal 107" in capsys.readouterr().out - - -def test_public_api_exports_expected_helpers(): - assert run_optimization is not None - assert check_optimization is not None - assert agent_spec_from_rlm is not None - assert build_proposer_for_rlm is not None - - -def test_apply_optimize_args_does_not_mutate_default_config(): - config = OptimizeConfig(executor_lm="before") - args = argparse.Namespace( - executor_lm="after", - executor_sub_lm=None, - executor_reasoning_effort=None, - executor_sub_lm_reasoning_effort=None, - proposer_lm=None, - proposer_sub_lm=None, - proposer_reasoning_effort=None, - proposer_sub_lm_reasoning_effort=None, - max_metric_calls=None, - minibatch_size=None, - concurrency=None, - max_iterations=None, - task_timeout=None, - proposer_timeout=None, - heartbeat_interval_seconds=None, - run_dir=None, - candidate_selection_strategy=None, - component_selection_strategy=None, - max_merge_attempts=None, - resume=False, - cache=False, - verbose_rlm=False, - debug_rlm=False, - merge_proposer=True, - ) - - updated = apply_optimize_args(config, args) - - assert config.executor_lm == "before" - assert updated.executor_lm == "after" - assert updated.merge_proposer is True - - -class _TimeoutProject(_Project): - async def evaluate_example(self, candidate, example, context): - await asyncio.sleep(1) - return RLMGepaExampleResult(score=1.0, feedback="ok", traces=[]) - - -class _ExampleTimeoutProject(_Project): - def task_timeout_for_example(self, example, default_timeout): - return 1 - - def task_resources_for_example(self, example): - return {"cpus": 2, "memory_mb": 4096} - - async def evaluate_example(self, candidate, example, context): - await asyncio.sleep(0.02) - return RLMGepaExampleResult( - score=1.0, - feedback=f"timeout={context.task_timeout} resources={dict(context.task_resources)}", - traces=[ - RunTrace( - status="completed", - model="test", - iterations=1, - max_iterations=1, - duration_ms=1, - ) - ], - example_id=str(example), - ) - - -class _ImmediateProject(_Project): - async def evaluate_example(self, candidate, example, context): - return RLMGepaExampleResult( - score=1.0, - feedback="", - traces=[{"status": "ok"}], - example_id=str(example), - ) - - -def test_adapter_progress_bar_updates_per_example(tmp_path: Path, monkeypatch): - import rlm_gepa.runtime.adapter as adapter_module - - events: list[tuple[str, object]] = [] - - class FakeTqdm: - def __init__(self, **kwargs): - events.append(("init", kwargs)) - - def set_postfix_str(self, value): - events.append(("postfix", value)) - - def update(self, value): - events.append(("update", value)) - - def close(self): - events.append(("close", None)) - - monkeypatch.setattr(adapter_module, "tqdm", FakeTqdm, raising=False) - adapter = RLMGepaAdapter( - project=_ImmediateProject(), - lm=_DummyLM(), - sub_lm=_DummyLM(), - max_iterations=1, - concurrency=2, - task_timeout=1, - output_dir=tmp_path, - run_id="run_test", - display_progress_bar=True, - ) - - batch = adapter.evaluate(["a", "b"], {"skill_instructions": "seed"}, capture_traces=True) - - assert batch.scores == [1.0, 1.0] - assert events[0] == ( - "init", - {"total": 2, "desc": " MB 0000 (2 tasks)", "leave": False, "unit": "task"}, - ) - assert [event for event in events if event[0] == "update"] == [("update", 1), ("update", 1)] - assert events[-1] == ("close", None) - - -def test_adapter_writes_eval_progress_events(tmp_path: Path): - adapter = RLMGepaAdapter( - project=_ImmediateProject(), - lm=_DummyLM(), - sub_lm=_DummyLM(), - max_iterations=1, - concurrency=2, - task_timeout=1, - output_dir=tmp_path, - run_id="run_test", - ) - - batch = adapter.evaluate(["a", "b"], {"skill_instructions": "seed"}, capture_traces=True) - - assert batch.scores == [1.0, 1.0] - events = [json.loads(line) for line in (tmp_path / "eval_progress.jsonl").read_text().splitlines()] - seen = [(event["example_id"], event["status"]) for event in events] - assert sorted(seen) == sorted( - [ - ("a", "started"), - ("b", "started"), - ("a", "completed"), - ("b", "completed"), - ] - ) - assert all(event["label"] == "MB 0000" for event in events) - assert events[-1]["score"] == 1.0 - - -def test_adapter_progress_bar_labels_valset(tmp_path: Path, monkeypatch): - import rlm_gepa.runtime.adapter as adapter_module - - events: list[tuple[str, object]] = [] - - class FakeTqdm: - def __init__(self, **kwargs): - events.append(("init", kwargs)) - - def set_postfix_str(self, value): - events.append(("postfix", value)) - - def update(self, value): - events.append(("update", value)) - - def close(self): - events.append(("close", None)) - - monkeypatch.setattr(adapter_module, "tqdm", FakeTqdm, raising=False) - adapter = RLMGepaAdapter( - project=_ImmediateProject(), - lm=_DummyLM(), - sub_lm=_DummyLM(), - max_iterations=1, - concurrency=2, - task_timeout=1, - output_dir=tmp_path, - run_id="run_test", - display_progress_bar=True, - valset_size=2, - ) - - adapter.evaluate(["a", "b"], {"skill_instructions": "seed"}, capture_traces=False) - - assert events[0] == ( - "init", - {"total": 2, "desc": " VALSET 0000 (2 tasks)", "leave": False, "unit": "task"}, - ) - - -def test_adapter_caps_task_trace_filename_length(tmp_path: Path): - long_label = "long_evaluation_kind_with_nested_context_segments_and_repeated_observation_windows_001" - adapter = RLMGepaAdapter( - project=_ImmediateProject(), - lm=_DummyLM(), - sub_lm=_DummyLM(), - max_iterations=1, - concurrency=1, - task_timeout=1, - output_dir=tmp_path, - run_id=( - "synthetic-domain-neutral-task-with-extended-run-identifier-and-many-" - "descriptive-segments-for-trace-filename-stress-20260522-010041" - ), - ) - - batch = adapter.evaluate( - [long_label], - {"skill_instructions": "seed"}, - capture_traces=False, - kind=long_label, - ) - - trace_files = list((tmp_path / "task_traces").glob("*.jsonl")) - assert batch.scores == [1.0] - assert len(trace_files) == 1 - assert len(trace_files[0].name) < 255 - row = json.loads(trace_files[0].read_text()) - assert row["event_id"].endswith("attempt_0000") - assert row["kind"] == long_label - - -def test_adapter_classifies_no_trace_repeat_batch_as_minibatch(tmp_path: Path, monkeypatch): - import rlm_gepa.runtime.adapter as adapter_module - - descriptions: list[str] = [] - - class FakeTqdm: - def __init__(self, **kwargs): - descriptions.append(kwargs["desc"]) - - def set_postfix_str(self, value): - pass - - def update(self, value): - pass - - def close(self): - pass - - monkeypatch.setattr(adapter_module, "tqdm", FakeTqdm, raising=False) - adapter = RLMGepaAdapter( - project=_ImmediateProject(), - lm=_DummyLM(), - sub_lm=_DummyLM(), - max_iterations=1, - concurrency=2, - task_timeout=1, - output_dir=tmp_path, - run_id="run_test", - display_progress_bar=True, - valset_size=2, + "status": "max_iterations", + "iterations": 5, + "max_iterations": 5, + "duration_ms": 5000, + }, + }, + { + "example_id": "project-timeout", + "status": "error", + "error": "RLM timeout at 300s", + "trace": { + "status": "error", + "iterations": 1, + "max_iterations": 5, + "duration_ms": 2000, + }, + }, + ] + ) ) - adapter.evaluate(["a", "b"], {"skill_instructions": "seed"}, capture_traces=True) - adapter.evaluate(["a", "b"], {"skill_instructions": "seed"}, capture_traces=False) - - assert descriptions == [" MB 0000 (2 tasks)", " MB 0001 (2 tasks)"] - - -def test_adapter_progress_bar_uses_reflective_context(tmp_path: Path, monkeypatch): - import rlm_gepa.runtime.adapter as adapter_module - - descriptions: list[str] = [] - - class FakeTqdm: - def __init__(self, **kwargs): - descriptions.append(kwargs["desc"]) - - def set_postfix_str(self, value): - pass - - def update(self, value): - pass - - def close(self): - pass + terminal = render_stats(tmp_path, table="all") + markdown = render_stats(tmp_path, table="all", output_format="markdown") - monkeypatch.setattr(adapter_module, "tqdm", FakeTqdm, raising=False) - adapter = RLMGepaAdapter( - project=_ImmediateProject(), - lm=_DummyLM(), - sub_lm=_DummyLM(), - max_iterations=1, - concurrency=2, - task_timeout=1, - output_dir=tmp_path, - run_id="run_test", - display_progress_bar=True, - valset_size=2, + expected = ( + "attempts=4, timeouts=2, max_iter_hits=1, latency p50=2.0s p90=5.0s p95=5.0s max=5.0s" ) - - adapter.set_reflective_progress_context(iteration=13, parent_idx=4, child_idx=7) - adapter.evaluate(["a", "b"], {"skill_instructions": "parent"}, capture_traces=True) - adapter.evaluate(["a", "b"], {"skill_instructions": "child"}, capture_traces=False) - adapter.evaluate(["c", "d"], {"skill_instructions": "child"}, capture_traces=False) - - assert descriptions == [ - " Iteration 13 Parent #4 Minibatch (2 tasks)", - " Iteration 13 Child #7 Minibatch (2 tasks)", - " Candidate #7 Valset (2 tasks)", - ] - - -class _ContextProject(_Project): - def __init__(self): - self.verbose_values: list[bool] = [] - self.debug_values: list[bool] = [] - - async def evaluate_example(self, candidate, example, context): - self.verbose_values.append(context.verbose_rlm) - self.debug_values.append(context.debug_rlm) - return RLMGepaExampleResult( - score=1.0, - feedback="", - traces=[{"status": "ok"}], - example_id=str(example), - ) + assert expected in terminal + assert expected in markdown -def test_adapter_propagates_rlm_logging_flags_to_every_example(tmp_path: Path): - project = _ContextProject() - adapter = RLMGepaAdapter( - project=project, - lm=_DummyLM(), - sub_lm=_DummyLM(), - max_iterations=1, - concurrency=2, - task_timeout=1, - output_dir=tmp_path, - run_id="run_test", - verbose_rlm=True, - debug_rlm=True, +def test_project_cli_check_with_dummy_lms(capsys): + config = OptimizeConfig( + executor_lm=_DummyLM(), + executor_sub_lm=_DummyLM(), + proposer_lm=_DummyLM(), + proposer_sub_lm=_DummyLM(), ) + status = run_project_cli(lambda: _Project(), config, argv=["optimize", "--check"]) - adapter.evaluate(["a", "b"], {"skill_instructions": "seed"}, capture_traces=False) - - assert project.verbose_values == [True, True] - assert project.debug_values == [True, True] + assert status == 0 + assert "check ok" in capsys.readouterr().out -class _TelemetryProject(_Project): - def __init__(self): - self.contexts: list[EvaluationContext] = [] +class _TimeoutProject(_Project): + cancelled = False async def evaluate_example(self, candidate, example, context): - self.contexts.append(context) - context.telemetry_context.write_span( - "test.case", - event_domain="test", - attributes={"example": example}, - ) - return RLMGepaExampleResult( - score=1.0, - feedback="", - traces=[{"status": "ok"}], - example_id=str(example), - ) - - -def test_adapter_threads_telemetry_context_and_persists_candidate_hash(tmp_path: Path): - project = _TelemetryProject() - telemetry_context = TelemetryContext( - sink=JsonlTelemetrySink(tmp_path / "telemetry" / "events.jsonl"), - trace_id="run_test", - run_id="run_test", - ) - adapter = RLMGepaAdapter( - project=project, - lm=_DummyLM(), - sub_lm=_DummyLM(), - max_iterations=1, - concurrency=2, - task_timeout=1, - output_dir=tmp_path, - run_id="run_test", - telemetry_context=telemetry_context, - ) - - adapter.evaluate( - [SimpleNamespace(task_id="example")], - {"skill_instructions": "seed"}, - capture_traces=False, - ) - - context = project.contexts[0].telemetry_context - assert context.run_id == "run_test" - assert context.eval_kind == "valset" - assert context.eval_idx == 0 - assert context.attempt_id == "attempt_0000" - assert context.example_id == "example" - assert context.candidate_id is None - assert context.candidate_hash.startswith("cand_sha256_") - events = [ - json.loads(line) - for line in (tmp_path / "telemetry" / "events.jsonl").read_text().splitlines() - ] - assert events[0]["attributes"]["rlm.candidate_hash"] == context.candidate_hash - trace_rows = [ - json.loads(line) - for line in ( - tmp_path / "task_traces" / "run_test_eval_valset_attempt_0000_valset.jsonl" - ) - .read_text() - .splitlines() - ] - assert trace_rows[0]["candidate_id"] is None - assert trace_rows[0]["candidate_hash"] == context.candidate_hash - assert trace_rows[0]["telemetry_ref"]["trace_id"] == context.trace_id + try: + await asyncio.sleep(1) + except asyncio.CancelledError: + self.cancelled = True + raise + return RLMGepaExampleResult(score=1.0, feedback="ok", traces=[]) class _FailingTelemetryProject(_Project): @@ -3247,43 +976,6 @@ async def evaluate_example(self, candidate, example, context): ) -def test_adapter_trace_rows_include_compact_failure_metadata(tmp_path: Path): - telemetry_context = TelemetryContext( - sink=JsonlTelemetrySink(tmp_path / "telemetry" / "events.jsonl"), - trace_id="run_test", - run_id="run_test", - ) - adapter = RLMGepaAdapter( - project=_FailingTelemetryProject(), - lm=_DummyLM(), - sub_lm=_DummyLM(), - max_iterations=1, - concurrency=1, - task_timeout=1, - output_dir=tmp_path, - run_id="run_test", - telemetry_context=telemetry_context, - ) - - adapter.evaluate(["example"], {"skill_instructions": "seed"}, capture_traces=False) - - trace_rows = [ - json.loads(line) - for line in ( - tmp_path / "task_traces" / "run_test_eval_valset_attempt_0000_valset.jsonl" - ) - .read_text() - .splitlines() - ] - row = trace_rows[0] - assert row["candidate_id"] is None - assert row["candidate_hash"].startswith("cand_sha256_") - assert row["failure_class"] == "host_tool_timeout_or_leak" - assert row["failure_reason"] == "tool timed out" - assert row["telemetry_ref"]["events_path"] == "telemetry/events.jsonl" - assert row["telemetry_ref"]["trace_id"].endswith(":0") - - def test_reflective_record_visible_to_gepa_includes_failure_metadata(tmp_path: Path): telemetry_context = TelemetryContext( sink=JsonlTelemetrySink(tmp_path / "telemetry" / "events.jsonl"), @@ -3444,69 +1136,10 @@ def test_adapter_reflective_records_include_structured_run_traces(tmp_path: Path assert task_row["traces"][0]["usage"]["sub"]["cost"] == 0.002 -def test_gepa_failure_metadata_includes_lm_truncation_fields(): - events = [ - { - "name": "rlm.action_generation.parse_error", - "status": {"code": "ERROR", "message": "parse failed"}, - "attributes": { - "failure.class": "model_output_truncated", - "lm.truncated": True, - "lm.truncation_reason": "max_tokens", - "lm.finish_reason": "length", - "lm.max_tokens": 50000, - "lm.output_tokens": 50000, - }, - } - ] - - metadata = _row_failure_metadata( - {"score": 0.0}, - events, - telemetry_context=None, - ) - - assert metadata["failure_class"] == "model_output_truncated" - assert metadata["truncated"] is True - assert metadata["truncation_reason"] == "max_tokens" - assert metadata["finish_reason"] == "length" - assert metadata["max_tokens"] == 50000 - assert metadata["output_tokens"] == 50000 - - -def test_gepa_failure_metadata_infers_truncation_from_structured_trace_finish_reason(): - metadata = _row_failure_metadata( - { - "score": 0.0, - "trace": { - "usage": { - "main": { - "output_tokens": 50000, - "max_tokens": 50000, - } - }, - "steps": [{"lm": {"finish_reason": "length"}}], - }, - }, - [], - telemetry_context=None, - ) - - assert metadata["failure_class"] == "model_output_truncated" - assert metadata["truncated"] is True - assert metadata["truncation_reason"] == "max_tokens" - assert metadata["finish_reason"] == "length" - assert metadata["max_tokens"] == 50000 - assert metadata["output_tokens"] == 50000 - - -def test_telemetry_classifier_uses_truncation_without_error_text(): - assert classify_failure({"score": 0.0, "truncated": True}, []) == "model_output_truncated" - - def test_adapter_enforces_per_example_timeout(tmp_path: Path): + project = _TimeoutProject() adapter = RLMGepaAdapter( - project=_TimeoutProject(), + project=project, lm=_DummyLM(), sub_lm=_DummyLM(), max_iterations=1, @@ -3519,50 +1152,10 @@ def test_adapter_enforces_per_example_timeout(tmp_path: Path): batch = adapter.evaluate(["example"], {"skill_instructions": "seed"}, capture_traces=True) assert batch.scores == [0.0] - assert batch.trajectories[0]["record"]["Feedback"] == "evaluation timeout at 0.01s" - - -def test_adapter_uses_project_timeout_and_resources_for_each_example(tmp_path: Path): - adapter = RLMGepaAdapter( - project=_ExampleTimeoutProject(), - lm=_DummyLM(), - sub_lm=_DummyLM(), - max_iterations=1, - concurrency=1, - task_timeout=0.01, - output_dir=tmp_path, - run_id="run_test", - ) - - batch = adapter.evaluate(["example"], {"skill_instructions": "seed"}, capture_traces=True) - - assert batch.scores == [1.0] - assert batch.trajectories[0]["record"]["Feedback"] == ( - "timeout=1 resources={'cpus': 2, 'memory_mb': 4096}" - ) - - -def test_adapter_prints_big_warning_for_evaluation_errors(tmp_path: Path, monkeypatch): - import rlm_gepa.runtime.adapter as adapter_module - - messages: list[str] = [] - monkeypatch.setattr(adapter_module, "progress_write", messages.append) - adapter = RLMGepaAdapter( - project=_ErrorProject(), - lm=_DummyLM(), - sub_lm=_DummyLM(), - max_iterations=1, - concurrency=1, - task_timeout=1, - output_dir=tmp_path, - run_id="run_test", - ) - - adapter.evaluate(["example"], {"skill_instructions": "seed"}) - - assert messages == [ - "⚠️ EVALUATION ERROR valset example: expected failure", - ] + assert project.cancelled + row = json.loads(next((tmp_path / "task_traces").glob("*.jsonl")).read_text()) + assert row["score"] == 0.0 + assert "timeout" in row["error"] class _ErrorProject(_Project): @@ -3582,7 +1175,9 @@ def test_resume_uses_unique_event_namespace_for_write_once_artifacts(tmp_path: P _run_dir, first_run_id = prepare_run_dir(_Project(), config, command="first") assert (run_dir / "telemetry").is_dir() (run_dir / "gepa_state.bin").write_bytes(b"checkpoint") - old_trace = run_dir / "task_traces" / f"{first_run_id}_eval_valset_attempt_0000_valset.jsonl" + old_trace = ( + run_dir / "task_traces" / f"{first_run_id}_eval_valset_attempt_0000_valset.jsonl" + ) old_trace.write_text("existing\n") resume_config = OptimizeConfig(run_dir=run_dir, resume=True) @@ -3604,142 +1199,15 @@ def test_resume_uses_unique_event_namespace_for_write_once_artifacts(tmp_path: P assert batch.scores == [0.0] assert old_trace.read_text() == "existing\n" - new_trace = run_dir / "task_traces" / f"{resume_run_id}_eval_valset_attempt_0000_valset.jsonl" - assert new_trace.exists() - - -def test_patch_merge_adapter_uses_patch_signature_and_persists_metadata( - tmp_path: Path, - monkeypatch, -): - import rlm_gepa.proposer.rlm as proposer_module - import rlm_gepa.runtime.adapter as adapter_module - - captured: dict[str, object] = {} - proposer_lm = _DummyLM() - proposer_sub_lm = _DummyLM() - - class FakePredictRLM: - def __init__(self, signature, *, lm, sub_lm, skills, max_iterations, verbose, debug): - captured.update( - { - "signature": signature, - "lm": lm, - "sub_lm": sub_lm, - "skills": skills, - "max_iterations": max_iterations, - "verbose": verbose, - "debug": debug, - } - ) - - async def acall(self, **kwargs): - captured["inputs"] = kwargs - return SimpleNamespace( - base_parent_id=10, - patch_summary="imported one clause", - selected_capability={ - "decision": "grafted", - "summary": "validate inputs before invoking tools", - "evidence_task_ids": ["train-a"], - "trigger": "task requires validating user-provided tool inputs before invocation", - "non_application_boundary": ( - "do not apply on base-win rows where direct tool invocation already succeeds" - ), - }, - patch_audit={ - "supported_source_win_ids": ["train-a"], - "guardrail_hazards": [], - "notes": "base lacks this validated-input facet", - }, - new_instructions="base plus patch", - trace=None, - trajectory=[], - ) - - monkeypatch.setattr(adapter_module, "PredictRLM", FakePredictRLM) - monkeypatch.setattr(adapter_module, "progress_write", lambda _message: None) - monkeypatch.setattr(proposer_module, "progress_write", lambda _message: None) - (tmp_path / "proposer_traces").mkdir() - paired_trace = tmp_path / "paired_patch.jsonl" - paired_trace.write_text("{}\n") - adapter = RLMGepaAdapter( - project=_Project(), - lm=_DummyLM(), - sub_lm=_DummyLM(), - max_iterations=1, - concurrency=1, - task_timeout=1, - output_dir=tmp_path, - run_id="run_test", - proposer_lm=proposer_lm, - proposer_sub_lm=proposer_sub_lm, - proposer_max_iterations=17, - ) - - new_text, metadata = adapter._rlm_propose_patch_merge_texts( - call_idx=4, - attempt_idx=2, - base_parent_id=10, - patch_source_parent_id=11, - base_parent_instructions="base", - patch_source_parent_instructions="source", - paired_disagreement_traces_file=SimpleNamespace(path=str(paired_trace)), - trace_task_ids=["train-a"], + new_trace = ( + run_dir / "task_traces" / f"{resume_run_id}_eval_valset_attempt_0000_valset.jsonl" ) - - assert new_text == "base plus patch" - assert metadata["patch_summary"] == "imported one clause" - assert metadata["selected_capability"]["decision"] == "grafted" - assert ( - metadata["selected_capability"]["trigger"] - == "task requires validating user-provided tool inputs before invocation" - ) - assert ( - metadata["selected_capability"]["non_application_boundary"] - == "do not apply on base-win rows where direct tool invocation already succeeds" - ) - assert metadata["base_instruction_chars"] == len("base") - assert metadata["new_instruction_chars"] == len("base plus patch") - assert metadata["instruction_char_delta"] == len("base plus patch") - len("base") - assert captured["signature"].input_fields.keys() >= { - "base_parent_id", - "base_parent_instructions", - "patch_source_parent_id", - "patch_source_parent_instructions", - "paired_disagreement_traces_file", - } - assert "selected_capability" in captured["signature"].output_fields - assert "patch_audit" in captured["signature"].output_fields - assert "behavioral_rules" not in captured["signature"].output_fields - assert "patch_merge_audit" not in captured["signature"].output_fields - assert "rejected_from_other" not in captured["signature"].output_fields - assert "imported_from_other" not in captured["signature"].output_fields - assert "common_ancestor_instructions" not in captured["inputs"] - assert captured["inputs"]["base_parent_id"] == 10 - assert captured["inputs"]["patch_source_parent_id"] == 11 - artifacts = list( - (tmp_path / "proposer_traces").glob("*_patch_from_cand_10_using_cand_11.json") - ) - assert len(artifacts) == 1 - payload = json.loads(artifacts[0].read_text()) - assert payload["kind"] == "patch_merge_proposer" - patch_output = payload["patch_output"] - _assert_valid_patch_output( - patch_output, - trace_task_ids=["train-a"], - base_instructions="base", - new_instructions="base plus patch", - ) - assert patch_output["patch_audit"]["supported_source_win_ids"] == ["train-a"] - assert "behavioral_rules" not in patch_output - assert "patch_merge_audit" not in patch_output - assert "rejected_from_other" not in patch_output + row = json.loads(new_trace.read_text()) + assert row["score"] == 0.0 + assert row["error"] == "expected failure" -def test_rlm_patch_merge_no_op_patch_persists_compact_audit( - tmp_path: Path, monkeypatch -): +def test_rlm_patch_merge_no_op_patch_persists_compact_audit(tmp_path: Path, monkeypatch): import rlm_gepa.proposer.rlm as proposer_module import rlm_gepa.runtime.adapter as adapter_module @@ -3808,91 +1276,17 @@ async def acall(self, **_kwargs): assert metadata["selected_capability"]["evidence_task_ids"] == [] assert metadata["instruction_char_delta"] == 0 assert metadata["patch_audit"]["supported_source_win_ids"] == [] - assert "duplicate" in metadata["patch_audit"]["notes"] artifacts = list( (tmp_path / "proposer_traces").glob("*_patch_from_cand_10_using_cand_11.json") ) assert len(artifacts) == 1 patch_output = json.loads(artifacts[0].read_text())["patch_output"] - _assert_valid_patch_output( - patch_output, - trace_task_ids=["train-a"], - base_instructions=base_instructions, - new_instructions=base_instructions, - ) + assert patch_output["new_instructions"] == base_instructions + assert patch_output["instruction_char_delta"] == 0 + assert patch_output["patch_audit"]["supported_source_win_ids"] == [] assert patch_output["selected_capability"]["decision"] == "no-op" -def test_patch_merge_prompt_contains_compact_grounding_invariants(): - instructions = build_merge_signature(_spec()).instructions - - assert "# Workflow" in instructions - assert "# Patch Contract" not in instructions - assert "one coherent missing capability family" in instructions - assert "necessary facets of the same behavior" in instructions - assert "do not import unrelated source-parent behaviors" in instructions - assert "return\n`new_instructions` unchanged" in instructions - assert "no task IDs, row labels, audit labels, or" in instructions - assert "provenance notes in `new_instructions`" in instructions - assert "base wins" in instructions - assert "both-success rows are preservation checks" in instructions - assert PatchMergeInstructionsGeneric.input_fields[ - "paired_disagreement_traces_file" - ].json_schema_extra["desc"] == ( - "JSONL file carrying train disagreement evidence for the two parents" - ) - assert "use those as primary behavioral evidence" in instructions - assert "Inspect failed rows alongside scores and feedback" in instructions - assert "tool-call inputs/outputs/errors" in instructions - assert "predict-call\n inputs/outputs/errors" in instructions - assert "LM finish reasons to understand why one parent\n won" in instructions - assert "`steps[*].output` can be shortened for display" in instructions - assert "prefer `steps[*].untruncated_output` if present" in instructions - assert 'failure_metadata.failure_class == "model_output_truncated"' in instructions - assert "generated LM answer was cut off or incomplete" in instructions - assert "shortened sandbox display output" in instructions - assert "Use repeated\n behavioral failure modes" in instructions - assert "do not overfit to one unusual disagreement row" in instructions - assert "Use available tools and `predict()` for focused evidence extraction" in instructions - assert "Helper `predict()` calls may extract evidence" in instructions - assert "you choose the final" in instructions - assert "Before editing, identify one patch-source-win cluster" in instructions - assert "state the exact" in instructions - assert "task-intent or observable trigger" in instructions - assert "state the" in instructions - assert "non-application boundary" in instructions - assert "grounded in base-win or both-success evidence" in instructions - assert "cannot state the trigger and boundary concretely" in instructions - assert "ProposerRunTrace" not in instructions - assert "archival" not in instructions - assert "token cost/cache accounting" not in instructions - assert "durations" not in instructions - assert "candidate_hash" not in instructions - assert "outer proposer" not in instructions - assert "helper output" not in instructions - assert "support-filter" not in instructions - assert "verification" not in instructions - - -def test_generic_proposer_prompt_contains_surgical_compression_invariants(): - instructions = build_proposer_signature(_spec()).instructions - - assert "spreadsheet formula" not in instructions.lower() - assert ImproveInstructionsGeneric.input_fields["traces_file"].json_schema_extra["desc"] == ( - "JSON file containing structured execution evidence for proposer review" - ) - assert "ProposerRunTrace" not in instructions - assert "archival" not in instructions - assert "token cost/cache accounting" not in instructions - assert "durations" not in instructions - assert "candidate_hash" not in instructions - assert "outer proposer" not in instructions - assert "helper output" not in instructions - assert "support-filter" not in instructions - assert "one coherent missing capability family" not in instructions - assert "patch-source" not in instructions - - def test_rlm_instruction_proposer_serializes_proposer_trace_records( tmp_path: Path, monkeypatch ): @@ -3916,8 +1310,6 @@ async def acall(self, **kwargs): monkeypatch.setattr(proposer_module, "PredictRLM", FakePredictRLM) monkeypatch.setattr(proposer_module, "progress_write", lambda _message: None) - monkeypatch.setattr(proposer_module, "install_rlm_log_stream", lambda _label: None) - monkeypatch.setattr(proposer_module, "restore_rlm_log_stream", lambda _stream: None) proposer = RLMInstructionProposer( spec=_spec(), lm=_DummyLM(), @@ -3996,14 +1388,13 @@ async def acall(self, **kwargs): serialized = captured["records"] assert isinstance(serialized, list) assert serialized[0]["Traces"][0]["steps"][0]["tool_calls"][0]["error"] == "boom" - assert serialized[0]["Traces"][0]["steps"][0]["predict_calls"][0]["calls"][0][ - "error" - ] == "predict boom" + assert ( + serialized[0]["Traces"][0]["steps"][0]["predict_calls"][0]["calls"][0]["error"] + == "predict boom" + ) serialized_text = json.dumps(serialized) assert "QUJDREVGRw==" not in serialized_text assert "data:image/png;base64," in serialized_text - assert "usage" not in serialized[0]["Traces"][0] - assert "duration_ms" not in serialized[0]["Traces"][0] assert "usage" not in serialized_text assert "duration_ms" not in serialized_text assert "cost" not in serialized_text @@ -4015,54 +1406,3 @@ async def acall(self, **kwargs): assert "trace_id" not in serialized_text assert "Trace Preview" not in serialized[0] assert "Generated Outputs" not in serialized[0] - - -def test_generic_proposer_output_fields_describe_compact_edits_and_preservation(): - output_fields = ImproveInstructionsGeneric.output_fields - - new_desc = output_fields["new_instructions"].json_schema_extra["desc"] - check_desc = output_fields["generalization_check"].json_schema_extra["desc"] - - assert "compact replacement or compression over appending" in new_desc - assert "Do not include audit labels or task IDs" in new_desc - assert "preserved solved behavior" in check_desc - - -def _assert_valid_patch_output( - patch_output: dict[str, object], - *, - trace_task_ids: list[str], - base_instructions: str, - new_instructions: str, -) -> None: - selected_capability = patch_output["selected_capability"] - assert isinstance(selected_capability, dict) - assert set(selected_capability) >= { - "decision", - "summary", - "evidence_task_ids", - "trigger", - "non_application_boundary", - } - assert set(selected_capability["evidence_task_ids"]) <= set(trace_task_ids) - assert isinstance(selected_capability["trigger"], str) - assert selected_capability["trigger"].strip() - assert isinstance(selected_capability["non_application_boundary"], str) - assert selected_capability["non_application_boundary"].strip() - patch_audit = patch_output["patch_audit"] - assert isinstance(patch_audit, dict) - assert set(patch_audit) >= { - "supported_source_win_ids", - "guardrail_hazards", - "notes", - } - assert patch_output["base_instruction_chars"] == len(base_instructions) - assert patch_output["new_instruction_chars"] == len(new_instructions) - assert patch_output["instruction_char_delta"] == len(new_instructions) - len( - base_instructions - ) - for task_id in trace_task_ids: - if len(task_id) >= 4: - assert task_id not in new_instructions - for audit_label in ("base_win", "patch_source_win", "both_success_guardrail"): - assert audit_label not in new_instructions diff --git a/tests/test_rlm_gepa_patch_merge_costs.py b/tests/test_rlm_gepa_patch_merge_costs.py index e7f00b1b..24f83c5c 100644 --- a/tests/test_rlm_gepa_patch_merge_costs.py +++ b/tests/test_rlm_gepa_patch_merge_costs.py @@ -31,7 +31,11 @@ class _Project(RLMGepaProject): agent_spec = AgentSpec( agent_type="test agent", use_cases=["case a", "case b"], - runtime_grounding_examples={"tools": ["tool()"], "env": ["sandbox"], "spec": ["protocol"]}, + runtime_grounding_examples={ + "tools": ["tool()"], + "env": ["sandbox"], + "spec": ["protocol"], + }, tool_signatures="tool() -> str", target_signature="input: str -> output: str", scoring_description="exact match", @@ -87,9 +91,13 @@ def _trace_with_proposer_usage() -> RunTrace: ) -def test_patch_merge_proposer_logs_merge_proposer_cost_roles(tmp_path: Path, monkeypatch): +def test_patch_merge_preserves_output_artifact_and_charges_main_and_sub_usage( + tmp_path: Path, monkeypatch +): import rlm_gepa.runtime.adapter as adapter_module + instructions = " Preserve base rules.\n\nApply the supported patch: café.\n" + class FakePredictRLM: def __init__(self, *_args, **_kwargs): pass @@ -112,7 +120,7 @@ async def acall(self, **_kwargs): "guardrail_hazards": [], "notes": "base lacks the selected facet", }, - new_instructions="patched instructions", + new_instructions=instructions, trace=_trace_with_proposer_usage(), trajectory=[], ) @@ -135,7 +143,7 @@ async def acall(self, **_kwargs): proposer_max_iterations=1, ) - adapter._rlm_propose_patch_merge_texts( + new_text, _metadata = adapter._rlm_propose_patch_merge_texts( call_idx=1, attempt_idx=0, base_parent_id=1, @@ -145,25 +153,29 @@ async def acall(self, **_kwargs): paired_disagreement_traces_file=SimpleNamespace(path=str(paired_trace)), trace_task_ids=["train-a"], ) - - cost_log = [json.loads(line) for line in (tmp_path / "cost_log.jsonl").read_text().splitlines()] + artifact = next( + (tmp_path / "proposer_traces").glob("*_patch_from_cand_1_using_cand_2.json") + ) + payload = json.loads(artifact.read_text()) + assert new_text == instructions + assert payload["new_instructions"] == instructions + assert payload["patch_output"]["new_instructions"] == instructions + assert payload["patch_output"]["instruction_char_delta"] == len(instructions) - len("base") + + cost_log = [ + json.loads(line) for line in (tmp_path / "cost_log.jsonl").read_text().splitlines() + ] assert [row["role"] for row in cost_log] == ["merge_proposer", "merge_proposer_sub_lm"] assert [row["cost_usd"] for row in cost_log] == [0.03, 0.02] rows = cost_rows(tmp_path) - assert any(row.get("scope") == "merge" and row.get("_category") for row in rows) - assert any( - row.get("scope") == " - proposer main" and row.get("total_cost") == "$0.03" - for row in rows - ) - assert any( - row.get("scope") == " - proposer sub" and row.get("total_cost") == "$0.02" - for row in rows - ) - assert not any(row.get("scope") == "patch-merge" for row in rows) + total = next(row for row in rows if row["scope"] == "TOTAL") + assert total["total_cost"] == "$0.05" + assert total["effective_cost"] == "$0.05" + assert total["repeat_cost"] == "$0.00" -def test_cost_rows_group_legacy_patch_merge_roles_under_merge_proposer(tmp_path: Path): +def test_legacy_patch_roles_preserve_both_costs_with_shared_operation_ids(tmp_path: Path): append_cost_rows( tmp_path / "cost_log.jsonl", [ @@ -196,13 +208,7 @@ def test_cost_rows_group_legacy_patch_merge_roles_under_merge_proposer(tmp_path: rows = cost_rows(tmp_path) - assert any(row.get("scope") == "merge" and row.get("_category") for row in rows) - assert any( - row.get("scope") == " - proposer main" and row.get("model") == "dummy-main" - for row in rows - ) - assert any( - row.get("scope") == " - proposer sub" and row.get("model") == "dummy-sub" - for row in rows - ) - assert not any(row.get("scope") == "patch-merge" for row in rows) + total = next(row for row in rows if row["scope"] == "TOTAL") + assert total["total_cost"] == "$0.03" + assert total["effective_cost"] == "$0.03" + assert total["repeat_cost"] == "$0.00" diff --git a/tests/test_rlm_skill_docs.py b/tests/test_rlm_skill_docs.py deleted file mode 100644 index c1d0b55d..00000000 --- a/tests/test_rlm_skill_docs.py +++ /dev/null @@ -1,28 +0,0 @@ -"""Regression checks for packaged RLM skill docs.""" - -from __future__ import annotations - -import re -import tomllib -from pathlib import Path - -ROOT = Path(__file__).resolve().parents[1] - - -def test_public_rlm_skill_version_snippets_match_package_version(): - package_version = tomllib.loads((ROOT / "pyproject.toml").read_text())["project"][ - "version" - ] - skill_text = (ROOT / ".agents/skills/rlm/SKILL.md").read_text() - - if package_version != "0.4.1": - return - - stale_snippets = [ - r"predict-rlm>=0\.3\.0", - r"predict-rlm\[[^\]]+\]>=0\.4\.0", - ] - for snippet in stale_snippets: - assert not re.search(snippet, skill_text), f"stale RLM skill snippet: {snippet}" - - diff --git a/tests/test_rlm_skills.py b/tests/test_rlm_skills.py index 3e8b468e..896dc4e6 100644 --- a/tests/test_rlm_skills.py +++ b/tests/test_rlm_skills.py @@ -1,120 +1,80 @@ -"""Tests for the Skill model and merge_skills.""" +"""Skill composition executes tools and refuses ambiguous resource ownership.""" +import dspy import pytest +from dspy.utils.dummies import DummyLM +from predict_rlm import PredictRLM from predict_rlm.rlm_skills import Skill, merge_skills -class TestSkill: - def test_minimal_skill(self): - skill = Skill(name="test") - assert skill.name == "test" - assert skill.instructions == "" - assert skill.packages == [] - assert skill.tools == {} +@pytest.mark.integration +def test_composed_skills_execute_together_in_the_rlm_loop(): + from predict_rlm.backends import JspiBackend - def test_full_skill(self): - def my_tool(x: str) -> str: - return x - - skill = Skill( - name="full", - instructions="Do the thing.", - packages=["pandas", "numpy"], - tools={"my_tool": my_tool}, - ) - assert skill.name == "full" - assert skill.instructions == "Do the thing." - assert skill.packages == ["pandas", "numpy"] - assert "my_tool" in skill.tools + def double(value: int) -> int: + return value * 2 + def label(value: int) -> str: + return f"result:{value}" -class TestMergeSkills: - def test_empty_list(self): - instructions, packages, modules, tools = merge_skills([]) - assert instructions == "" - assert packages == [] - assert tools == {} - - def test_single_skill(self): - skill = Skill( - name="pdf", - instructions="Use pdfplumber.", - packages=["pdfplumber"], + interpreter = JspiBackend(preinstall_packages=False) + try: + rlm = PredictRLM( + "query -> answer", + interpreter=interpreter, + skills=[ + Skill(name="arithmetic", tools={"double": double}), + Skill(name="report", tools={"label": label}), + ], + max_iterations=1, + ) + with dspy.context( + lm=DummyLM( + [ + { + "reasoning": "compose both skills", + "code": "value = await double(21)\nSUBMIT(answer=await label(value))", + } + ] + ) + ): + result = rlm(query="compute") + assert result.answer == "result:42" + finally: + interpreter.shutdown() + + +def test_duplicate_skill_tool_names_are_rejected(): + def tool(): + return "value" + + with pytest.raises(ValueError, match="Tool name conflict.*shared"): + merge_skills( + [ + Skill(name="first", tools={"shared": tool}), + Skill(name="second", tools={"shared": tool}), + ] ) - instructions, packages, modules, tools = merge_skills([skill]) - assert "Skill: pdf" in instructions - assert "Use pdfplumber." in instructions - assert packages == ["pdfplumber"] - assert tools == {} - - def test_multiple_skills_merge_instructions(self): - s1 = Skill(name="a", instructions="Do A.") - s2 = Skill(name="b", instructions="Do B.") - instructions, _, _, _ = merge_skills([s1, s2]) - assert "Skill: a" in instructions - assert "Skill: b" in instructions - assert "Do A." in instructions - assert "Do B." in instructions - - def test_package_dedup(self): - s1 = Skill(name="a", packages=["pandas", "numpy"]) - s2 = Skill(name="b", packages=["numpy", "scipy"]) - _, packages, _, _ = merge_skills([s1, s2]) - assert packages == ["pandas", "numpy", "scipy"] - - def test_tool_merge(self): - def tool_a(): - pass - - def tool_b(): - pass - - s1 = Skill(name="a", tools={"tool_a": tool_a}) - s2 = Skill(name="b", tools={"tool_b": tool_b}) - _, _, _, tools = merge_skills([s1, s2]) - assert set(tools.keys()) == {"tool_a", "tool_b"} - - def test_tool_name_conflict_raises(self): - def tool_x(): - pass - - s1 = Skill(name="a", tools={"shared": tool_x}) - s2 = Skill(name="b", tools={"shared": tool_x}) - with pytest.raises(ValueError, match="Tool name conflict.*shared"): - merge_skills([s1, s2]) - def test_empty_instructions_skipped(self): - s1 = Skill(name="a", instructions="") - s2 = Skill(name="b", instructions=" ") - s3 = Skill(name="c", instructions="Real instructions.") - instructions, _, _, _ = merge_skills([s1, s2, s3]) - assert "Skill: a" not in instructions - assert "Skill: b" not in instructions - assert "Skill: c" in instructions - def test_module_merge(self): - s1 = Skill(name="a", modules={"formula_eval": "/path/to/formula_eval.py"}) - s2 = Skill(name="b", modules={"chart_helper": "/path/to/chart_helper.py"}) - _, _, modules, _ = merge_skills([s1, s2]) - assert modules == { - "formula_eval": "/path/to/formula_eval.py", - "chart_helper": "/path/to/chart_helper.py", - } +def test_duplicate_skill_module_names_are_rejected(): + with pytest.raises(ValueError, match="Module name conflict.*shared_mod"): + merge_skills( + [ + Skill(name="first", modules={"shared_mod": "/path/a.py"}), + Skill(name="second", modules={"shared_mod": "/path/b.py"}), + ] + ) - def test_module_name_conflict_raises(self): - s1 = Skill(name="a", modules={"shared_mod": "/path/a.py"}) - s2 = Skill(name="b", modules={"shared_mod": "/path/b.py"}) - with pytest.raises(ValueError, match="Module name conflict.*shared_mod"): - merge_skills([s1, s2]) - def test_single_skill_with_modules(self): - s = Skill(name="spreadsheet", modules={"formula_eval": "/path/formula_eval.py"}) - _, _, modules, _ = merge_skills([s]) - assert modules == {"formula_eval": "/path/formula_eval.py"} +def test_skill_tool_cannot_silently_replace_user_tool(): + def tool(): + return "value" - def test_empty_modules_skipped(self): - s1 = Skill(name="a", modules={}) - s2 = Skill(name="b", modules={"mod": "/path/mod.py"}) - _, _, modules, _ = merge_skills([s1, s2]) - assert modules == {"mod": "/path/mod.py"} + with pytest.raises(ValueError, match="Tool name conflict.*shared"): + PredictRLM( + "query -> answer", + skills=[Skill(name="skill", tools={"shared": tool})], + tools={"shared": tool}, + ) diff --git a/tests/test_runtime_hooks.py b/tests/test_runtime_hooks.py index a2ecc474..179a405d 100644 --- a/tests/test_runtime_hooks.py +++ b/tests/test_runtime_hooks.py @@ -76,13 +76,9 @@ def runner(tmp_path): proc.close() -def test_runtime_hooks_are_opt_in(runner: LocalRunner, tmp_path: Path): - path = tmp_path / "no-hook.txt" - result = runner.request("execute", {"code": f"open({str(path)!r}, 'w').close()"}) - assert result["result"]["output"] == "" - - -def test_runtime_hooks_emit_function_events(runner: LocalRunner, tmp_path: Path): +def test_runtime_hook_registration_emits_user_events_and_can_be_cleared( + runner: LocalRunner, tmp_path: Path +): path = tmp_path / "hooked.txt" registered = runner.request( "register_runtime_hooks", @@ -121,6 +117,15 @@ def test_runtime_hooks_emit_function_events(runner: LocalRunner, tmp_path: Path) ] assert events[0]["phase"] == "before" assert events[1]["phase"] == "after" + assert path.read_text(encoding="utf-8") == "hello" + + runner.request("register_runtime_hooks", {"hooks": []}) + execute = runner.request( + "execute", + {"code": f"Path({str(path)!r}).write_text('after reset')"}, + ) + assert "method" not in execute + assert path.read_text(encoding="utf-8") == "after reset" def test_runtime_hooks_do_not_emit_internal_capture_file_events(runner: LocalRunner): @@ -143,19 +148,6 @@ def test_runtime_hooks_emit_error_events(runner: LocalRunner): assert event["method"] == "runtime_hook_event" assert event["params"]["target"] == "builtins.open" assert event["params"]["phase"] == "error" - - -def test_runtime_hooks_reset_on_reregister(runner: LocalRunner, tmp_path: Path): - runner.request( - "register_runtime_hooks", - {"hooks": [{"target": "pathlib.Path.write_text", "phases": ["before"]}]}, - ) - # Re-register with empty set should clear hooks. - runner.request("register_runtime_hooks", {"hooks": []}) - path = tmp_path / "after-reset.txt" - execute = runner.request( - "execute", - {"code": f"from pathlib import Path\nPath({str(path)!r}).write_text('x')"}, - ) - assert "method" not in execute - assert execute["result"]["output"] == "" + response = json.loads(runner.proc.stdout.readline()) + assert response["error"]["data"]["type"] == "FileNotFoundError" + assert response["error"]["message"] == event["params"]["error"] diff --git a/tests/test_sandbox_backend_benchmark.py b/tests/test_sandbox_backend_benchmark.py deleted file mode 100644 index 0d392534..00000000 --- a/tests/test_sandbox_backend_benchmark.py +++ /dev/null @@ -1,209 +0,0 @@ -from __future__ import annotations - -import importlib.util -import sys -import time -from pathlib import Path - -import pytest - -SCRIPT = ( - Path(__file__).parents[1] - / "scripts" - / "benchmarks" - / "sandbox_backend_benchmark.py" -) - - -def load_benchmark_module(): - spec = importlib.util.spec_from_file_location("sandbox_backend_benchmark", SCRIPT) - assert spec is not None - module = importlib.util.module_from_spec(spec) - assert spec.loader is not None - sys.modules[spec.name] = module - spec.loader.exec_module(module) - return module - - -def test_aggregate_stats_handles_empty_and_percentiles(): - bench = load_benchmark_module() - - empty = bench.aggregate_stats([]) - assert empty.mean_seconds is None - assert empty.p50_seconds is None - assert empty.p95_seconds is None - - stats = bench.aggregate_stats([0.5, 0.1, 0.3, 0.2, 0.4]) - assert stats.mean_seconds == pytest.approx(0.3) - assert stats.p50_seconds == pytest.approx(0.3) - assert stats.p95_seconds == pytest.approx(0.5) - - -def test_scenario_selection_filters_by_backend(): - bench = load_benchmark_module() - - assert bench.selected_backends("all") == ["jspi", "sbx", "sbx-pool"] - assert bench.applicable_scenarios("jspi", None) == [ - "warm_tiny", - "warm_fib_recursive", - "warm_fib_iterative", - ] - assert bench.applicable_scenarios("jspi", ["startup_execute_shutdown"]) == [ - "startup_execute_shutdown" - ] - assert bench.applicable_scenarios("sbx", ["pool_tiny", "warm_tiny"]) == ["warm_tiny"] - assert bench.applicable_scenarios("sbx-pool", None) == ["pool_tiny"] - assert bench.applicable_scenarios("sbx-pool", ["warm_tiny", "pool_tiny"]) == ["pool_tiny"] - assert bench.applicable_scenarios( - "sbx-pool", ["pool_replenish_tiny", "pool_replenish_fib_iterative"] - ) == ["pool_replenish_tiny", "pool_replenish_fib_iterative"] - - -def test_parser_defaults_and_sbx_pool_size_default(): - bench = load_benchmark_module() - - args = bench.make_parser().parse_args([]) - bench.validate_args(args) - - assert args.backend == "all" - assert args.tasks is None - assert args.scenario is None - assert args.concurrency == 1 - assert args.sbx_pool_size is None - assert args.sbx_buffer_size is None - - assert bench.task_count_for_scenario("startup_execute_shutdown", args.tasks) == 5 - assert bench.task_count_for_scenario("warm_tiny", args.tasks) == 100 - assert bench.task_count_for_scenario("warm_fib_recursive", args.tasks) == 25 - assert bench.task_count_for_scenario("pool_tiny", args.tasks) == 100 - assert bench.task_count_for_scenario("pool_replenish_tiny", args.tasks) == 100 - assert bench.task_count_for_scenario("pool_replenish_fib_iterative", args.tasks) == 25 - - -def test_validate_args_rejects_negative_buffer_size(): - bench = load_benchmark_module() - - args = bench.make_parser().parse_args(["--sbx-buffer-size", "-1"]) - - with pytest.raises(SystemExit, match="--sbx-buffer-size must be at least 0"): - bench.validate_args(args) - - -def test_sbx_buffer_size_defaults_to_pool_size(): - bench = load_benchmark_module() - - assert bench.sbx_buffer_size_for_pool(pool_size=5, requested_buffer_size=None) == 5 - assert bench.sbx_buffer_size_for_pool(pool_size=5, requested_buffer_size=2) == 2 - - -def test_parser_rejects_old_surge_flag(): - bench = load_benchmark_module() - - with pytest.raises(SystemExit): - bench.make_parser().parse_args(["--sbx-surge", "1"]) - - -def test_tasks_override_applies_to_all_scenarios(): - bench = load_benchmark_module() - - args = bench.make_parser().parse_args(["--tasks", "7"]) - bench.validate_args(args) - - for scenario in bench.SCENARIOS: - assert bench.task_count_for_scenario(scenario, args.tasks) == 7 - - -def test_build_result_counts_failures_without_negative_ok(): - bench = load_benchmark_module() - - result = bench.build_result( - backend="sbx", - scenario="warm_tiny", - tasks=100, - wall_seconds=0.2, - durations=[], - failures=["startup failed"], - ) - - assert result.ok == 0 - assert result.failed == 1 - assert result.tasks_per_second == 0.0 - - -def test_pool_replenishment_starts_when_sandbox_is_acquired(monkeypatch): - bench = load_benchmark_module() - replacement_started = bench.threading.Event() - execute_entered = bench.threading.Event() - created: list["FakeInterpreter"] = [] - - class FakeInterpreter: - def __init__(self, index: int) -> None: - self.index = index - - def prewarm(self) -> None: - created.append(self) - if self.index == 2: - replacement_started.set() - - def shutdown(self) -> None: - pass - - def make_fake_interpreter(prefix: str, index: int) -> FakeInterpreter: - return FakeInterpreter(index) - - def execute_with_replacement_overlap(interpreter: FakeInterpreter, code: str) -> None: - assert interpreter.index == 0 - assert code == bench.TINY_CODE - execute_entered.set() - assert replacement_started.wait(timeout=1.0) - - monkeypatch.setattr(bench, "make_sbx_pool_interpreter", make_fake_interpreter) - monkeypatch.setattr(bench, "execute_code", execute_with_replacement_overlap) - - result = bench.run_pool_replenish( - "pool_replenish_tiny", - tasks=1, - concurrency=1, - pool_size=1, - buffer_size=1, - fail_fast=True, - ) - - assert result.failed == 0 - assert [interpreter.index for interpreter in created] == [0, 1, 2] - - -def test_pool_replenishment_wall_time_excludes_slow_retirement(monkeypatch): - bench = load_benchmark_module() - slow_shutdown_seconds = 0.2 - - class FakeInterpreter: - def __init__(self, index: int) -> None: - self.index = index - - def prewarm(self) -> None: - pass - - def shutdown(self) -> None: - time.sleep(slow_shutdown_seconds) - - def make_fake_interpreter(prefix: str, index: int) -> FakeInterpreter: - return FakeInterpreter(index) - - def execute_fast(interpreter: FakeInterpreter, code: str) -> None: - assert code == bench.TINY_CODE - - monkeypatch.setattr(bench, "make_sbx_pool_interpreter", make_fake_interpreter) - monkeypatch.setattr(bench, "execute_code", execute_fast) - - result = bench.run_pool_replenish( - "pool_replenish_tiny", - tasks=1, - concurrency=1, - pool_size=1, - buffer_size=1, - fail_fast=True, - ) - - assert result.failed == 0 - assert result.wall_seconds < slow_shutdown_seconds / 2 diff --git a/tests/test_sbx_interpreter.py b/tests/test_sbx_interpreter.py index 597d3907..481f8206 100644 --- a/tests/test_sbx_interpreter.py +++ b/tests/test_sbx_interpreter.py @@ -3,12 +3,9 @@ from __future__ import annotations import asyncio -import base64 import json -import logging import os import queue -import secrets import select import shutil import socket @@ -16,28 +13,19 @@ import sys import threading import time -from collections import UserDict -from dataclasses import dataclass from pathlib import Path from types import SimpleNamespace from typing import Annotated -from unittest.mock import AsyncMock, MagicMock, patch +from unittest.mock import AsyncMock, MagicMock import pytest -from pydantic import ValidationError pytest.importorskip("websockets") # SBX/supervisor backend requires the [sbx] extra -pytestmark = pytest.mark.sbx -from dspy.primitives.code_interpreter import CodeInterpreterError, FinalOutput # noqa: E402 +pytestmark = pytest.mark.sbx -from predict_rlm.backends import ( # noqa: E402 - DEFAULT_SBX_TEMPLATE, - SbxBackend, - SbxConfig, - SbxPool, -) +from predict_rlm.backends import SbxBackend, SbxConfig # noqa: E402 from predict_rlm.backends.base import ( # noqa: E402 BackendExecutionGate, SandboxExecutionError, @@ -53,9 +41,15 @@ HostDirectoryMount, UnsupportedOperationError, ) -from predict_rlm.workspace import DirectWorkspaceMount # noqa: E402 -PAYLOAD_PATH = Path(__file__).parents[1] / "src" / "predict_rlm" / "backends" / "supervisor" / "_payload.py" +PAYLOAD_PATH = ( + Path(__file__).parents[1] + / "src" + / "predict_rlm" + / "backends" + / "supervisor" + / "_payload.py" +) def _drain_available_pipe_text(pipe) -> str: @@ -147,96 +141,6 @@ def close(self) -> None: self.proc.wait(timeout=5) -class SequentialActions: - def __init__(self, *actions: SimpleNamespace) -> None: - self.actions = list(actions) - self.calls: list[dict] = [] - - def __call__(self, **kwargs): - self.calls.append(kwargs) - assert self.actions, "PredictRLM requested more actions than the test provided" - return self.actions.pop(0) - - -class PredictionStub: - def __init__(self, answer: str) -> None: - self.answer = answer - - def keys(self) -> list[str]: - return ["answer"] - - def __getitem__(self, key: str) -> str: - return getattr(self, key) - - -def assert_predict_rlm_recovers_after_user_exceptions_and_tools_still_work( - pool: SbxPool, -) -> None: - from predict_rlm import PredictRLM - - async def host_echo(text: str) -> dict: - await asyncio.sleep(0) - return {"text": f"echo:{text}"} - - actions = SequentialActions( - SimpleNamespace( - reasoning="raise a normal exception from inside a loop", - code=( - "for idx in range(3):\n" - " print('loop idx', idx)\n" - " if idx == 1:\n" - " raise ValueError(f'bad loop idx {idx}')\n" - ), - ), - SimpleNamespace( - reasoning="exercise a missing variable path", - code="print(missing_recovery_variable)\n", - ), - SimpleNamespace( - reasoning="prove host callbacks still work after ordinary exceptions", - code=( - "prediction = await predict('question: str -> answer: str', " - "question='after exceptions')\n" - "echoed = await host_echo(prediction['answer'])\n" - "SUBMIT(answer=echoed['text'])" - ), - ), - ) - - mock_lm = MagicMock() - mock_predictor = MagicMock() - mock_predictor.acall = AsyncMock(return_value=PredictionStub("tool-ok")) - rlm = PredictRLM( - "prompt -> answer", - sub_lm=mock_lm, - max_iterations=3, - tools={"host_echo": host_echo}, - sandbox_backend="sbx", - sbx_pool=pool, - ) - rlm.generate_action = actions - - with patch("predict_rlm.predict_rlm.dspy.Predict", return_value=mock_predictor): - prediction = rlm(prompt="exercise exception recovery") - - assert prediction.answer == "echo:tool-ok" - assert [call["iteration"] for call in actions.calls] == ["1/3", "2/3", "3/3"] - assert mock_predictor.acall.await_count == 1 - assert mock_predictor.acall.await_args.kwargs["question"] == "after exceptions" - assert len(prediction.trace.steps) == 3 - - value_error_step, name_error_step, final_step = prediction.trace.steps - assert "for idx in range(3)" in value_error_step.code - assert "raise ValueError" in value_error_step.code - assert "[Error]" in value_error_step.untruncated_output - assert "ValueError" in value_error_step.untruncated_output - assert "bad loop idx 1" in value_error_step.untruncated_output - assert "[Error]" in name_error_step.untruncated_output - assert "NameError" in name_error_step.untruncated_output - assert "missing_recovery_variable" in name_error_step.untruncated_output - assert final_step.output == "FINAL: {'answer': 'echo:tool-ok'}" - - @pytest.fixture def runner(tmp_path): proc = LocalRunner(tmp_path) @@ -263,29 +167,6 @@ def __reduce__(self): assert snapshot["lost_globals"] == ["native_model"] assert NativeLike.reduce_called is False - def test_snapshot_preserves_safe_dataclass_and_mapping_values(self): - @dataclass - class RunSummary: - name: str - scores: list[int] - output_path: Path - - snapshot = _pickleable_globals_snapshot({ - "summary": RunSummary("mjcf", [1, 2], Path("/app/model.xml")), - "config": UserDict({"threshold": 0.6, "labels": ("fast", "exact")}), - }) - - assert snapshot["lost_globals"] == [] - assert snapshot["restored_globals"] == ["config", "summary"] - assert snapshot["globals"] == { - "summary": { - "name": "mjcf", - "scores": [1, 2], - "output_path": Path("/app/model.xml"), - }, - "config": {"threshold": 0.6, "labels": ("fast", "exact")}, - } - def test_snapshot_crosses_runner_queue_after_hard_timeout(self, runner: LocalRunner): runner.request( "execute", @@ -326,9 +207,7 @@ def test_snapshot_crosses_runner_queue_after_hard_timeout(self, runner: LocalRun class TestPythonRunnerProtocol: - def test_user_subprocess_stdin_is_isolated_from_runner_protocol( - self, runner: LocalRunner - ): + def test_user_subprocess_stdin_is_isolated_from_runner_protocol(self, runner: LocalRunner): code = ( "import subprocess, sys\n" "subprocess.run(\n" @@ -345,25 +224,6 @@ def test_user_subprocess_stdin_is_isolated_from_runner_protocol( assert result["result"]["output"].strip() == "" assert followup["result"]["output"].strip() == "123" - def test_reset_clears_globals_but_runner_process_survives(self, runner: LocalRunner): - before = runner.request("execute", {"code": "x = 40\nprint('ready')"}) - reset = runner.request("reset") - after = runner.request("execute", {"code": "print('x' in globals())"}) - - assert before["result"]["output"].strip() == "ready" - assert reset["result"] == {} - assert after["result"]["output"].strip() == "False" - - def test_submit_returns_final_payload(self, runner: LocalRunner): - runner.request( - "register_output_fields", - {"fields": [{"name": "answer", "annotation": "str"}]}, - ) - - result = runner.request("execute", {"code": "SUBMIT(answer='done')"}) - - assert result["result"]["final"] == {"answer": "done"} - def test_deferred_submit_preserves_background_service_for_confirmation( self, runner: LocalRunner ): @@ -443,57 +303,6 @@ def test_deferred_submit_preserves_background_service_for_confirmation( assert final["result"] == {"final": {"answer": "confirmed"}} assert post_final_result == 0 - def test_predict_image_data_url_round_trips_to_host_tool(self, runner: LocalRunner): - png_bytes = b"\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR" - runner.request("register_tools", {"tools": ["predict"]}) - - tool_call = runner.request( - "execute", - { - "code": ( - "import base64\n" - f"open('/sandbox/image.png', 'wb').write({png_bytes!r})\n" - "image_bytes = open('/sandbox/image.png', 'rb').read()\n" - "data_url = 'data:image/png;base64,' + base64.b64encode(image_bytes).decode()\n" - "result = await predict(\n" - " 'image: dspy.Image, question: str -> visible_text: str',\n" - " image=data_url,\n" - " question='What text is visible?',\n" - ")\n" - "print(result.visible_text)\n" - ) - }, - ) - - assert tool_call["method"] == "tool_call" - assert tool_call["params"]["name"] == "predict" - assert tool_call["params"]["args"] == [ - "image: dspy.Image, question: str -> visible_text: str" - ] - assert tool_call["params"]["kwargs"]["question"] == "What text is visible?" - image = tool_call["params"]["kwargs"]["image"] - assert image.startswith("data:image/png;base64,") - assert base64.b64decode(image.removeprefix("data:image/png;base64,")) == png_bytes - - assert runner.proc.stdin is not None - assert runner.proc.stdout is not None - runner.proc.stdin.write( - json.dumps( - { - "jsonrpc": "2.0", - "id": tool_call["id"], - "result": {"type": "json", "value": '{"visible_text": "hello"}'}, - } - ) - + "\n" - ) - runner.proc.stdin.flush() - response = json.loads(runner.proc.stdout.readline()) - - assert response["result"]["output"] == "hello\n" - assert runner.request("shutdown")["result"] == {"shutdown": True} - runner.proc.wait(timeout=5) - def test_kernel_result_waits_for_tool_reader_handoff(self, monkeypatch): from predict_rlm.backends.supervisor import _payload @@ -543,16 +352,8 @@ def reader_loop() -> None: assert not reader.is_alive() assert result_queue.get_nowait() == {"ok": True} - # TEMPORARY SKIP: drives many concurrent host-tool calls over a local _payload.py - # subprocess and reads replies on tight (3s) select timeouts. On loaded CI runners - # the supervisor is slow to respond and the reader/teardown intermittently times out. - # Not a product bug -- a CI-timing-sensitive subprocess test. Re-enable with - # CI-tolerant helper timeouts (follow-up). - # local-only: flaky on CI (subprocess-timing-sensitive); re-enable with CI-tolerant timeouts @pytest.mark.local - def test_stale_concurrent_tool_calls_do_not_poison_later_execute( - self, runner: LocalRunner - ): + def test_stale_concurrent_tool_calls_do_not_poison_later_execute(self, runner: LocalRunner): runner.request("register_tools", {"tools": ["predict"]}) first_execute_id = runner.send( "execute", @@ -617,70 +418,6 @@ def test_stale_concurrent_tool_calls_do_not_poison_later_execute( assert saw_followup_tool_call - def test_syntax_error_uses_json_rpc_error(self, runner: LocalRunner): - result = runner.request("execute", {"code": "for"}) - - assert result["error"]["data"]["type"] == "SyntaxError" - - def test_timeout_returns_structured_result_and_runner_survives( - self, runner: LocalRunner - ): - result = runner.request( - "execute", - { - "code": ( - "import sys\n" - "print('before timeout')\n" - "print('stderr before timeout', file=sys.stderr)\n" - "partial_timeout_state = 42\n" - "while True:\n" - " pass\n" - ), - "execution_timeout_seconds": 0.1, - }, - ) - followup = runner.request( - "execute", - {"code": "print('partial_timeout_state' in globals())\nprint('still alive')"}, - ) - - assert result["result"] == { - "timeout": {"seconds": 0.1}, - "stdout": "before timeout\n", - "stderr": "stderr before timeout\n", - "state": { - "preserved": True, - "source": "live_kernel", - "scope": "full_live", - }, - } - assert followup["result"]["output"] == "True\nstill alive\n" - - def test_execute_captures_child_process_output_with_timeout( - self, runner: LocalRunner - ): - result = runner.request( - "execute", - { - "code": ( - "import subprocess, sys\n" - "subprocess.run([\n" - " sys.executable,\n" - " '-c',\n" - " \"import sys; print('child stdout'); " - "print('child stderr', file=sys.stderr)\",\n" - "])\n" - ), - "execution_timeout_seconds": 2, - }, - ) - followup = runner.request("execute", {"code": "print('runner still usable')"}) - leaked_stderr = _drain_available_pipe_text(runner.proc.stderr) - - assert result["result"]["output"] == "child stdout\nchild stderr\n" - assert followup["result"]["output"] == "runner still usable\n" - assert leaked_stderr == "" - def test_timeout_preserves_child_process_output_and_runner_survives( self, runner: LocalRunner ): @@ -694,7 +431,7 @@ def test_timeout_preserves_child_process_output_and_runner_survives( " '-c',\n" " \"import sys, time; print('child before timeout'); " "print('child err before timeout', file=sys.stderr); " - "sys.stdout.flush(); sys.stderr.flush(); time.sleep(30)\",\n" + 'sys.stdout.flush(); sys.stderr.flush(); time.sleep(30)",\n' "])\n" ), "execution_timeout_seconds": 0.2, @@ -723,9 +460,7 @@ def test_unbounded_execute_runner_exit_returns_error_and_supervisor_survives( assert "execution runner exited without a result" in result["error"]["message"] assert followup["result"]["output"] == "supervisor survived runner exit\n" - def test_timeout_is_not_swallowed_by_user_exception_handler( - self, runner: LocalRunner - ): + def test_timeout_is_not_swallowed_by_user_exception_handler(self, runner: LocalRunner): result = runner.request( "execute", { @@ -754,68 +489,6 @@ def test_timeout_is_not_swallowed_by_user_exception_handler( } assert followup["result"]["output"] == "still alive\n" - def test_timeout_cancels_pending_async_work(self, runner: LocalRunner): - result = runner.request( - "execute", - { - "code": ( - "import asyncio\n" - "async def mutate_late():\n" - " await asyncio.sleep(1)\n" - " globals()['late_mutation'] = 'leaked'\n" - "await mutate_late()\n" - ), - "execution_timeout_seconds": 0.1, - }, - ) - followup = runner.request( - "execute", - { - "code": ( - "import asyncio\n" - "await asyncio.sleep(0.3)\n" - "print('late_mutation' in globals())" - ) - }, - ) - - assert result["result"]["timeout"] == {"seconds": 0.1} - assert followup["result"]["output"] == "False\n" - - def test_file_helpers_preserve_virtual_paths(self, runner: LocalRunner, tmp_path: Path): - source = tmp_path / "input.txt" - source.write_text("hello", encoding="utf-8") - out = tmp_path / "out.txt" - - runner.request( - "mount_file", - {"host_path": str(source), "virtual_path": "/sandbox/input/source/input.txt"}, - ) - runner.request("mkdir_p", {"path": "/sandbox/output/result"}) - runner.request( - "execute", - { - "code": ( - "with open('/sandbox/input/source/input.txt', encoding='utf-8') as f:\n" - " text = f.read()\n" - "with open('/sandbox/output/result/output.txt', 'w', encoding='utf-8') as f:\n" - " f.write(text + ' world')" - ) - }, - ) - - files = runner.request("list_dir", {"path": "/sandbox/output/result"}) - runner.request( - "sync_file", - { - "virtual_path": "/sandbox/output/result/output.txt", - "host_path": str(out), - }, - ) - - assert files["result"]["files"] == ["/sandbox/output/result/output.txt"] - assert out.read_text(encoding="utf-8") == "hello world" - def test_pathlib_path_remains_a_type(self, runner: LocalRunner): result = runner.request( "execute", @@ -830,143 +503,6 @@ def test_pathlib_path_remains_a_type(self, runner: LocalRunner): assert result["result"]["output"] == "True\nFalse\n" - def test_windows311_visual_predict_path_handles_pillow_style_path_checks( - self, runner: LocalRunner - ): - runner.request("register_tools", {"tools": ["predict"]}) - tool_call = runner.request( - "execute", - { - "code": ( - "import base64, pathlib\n" - "ppm = b'P6\\n1 1\\n255\\n' + bytes([255, 255, 255])\n" - "pathlib.Path('/tmp/win311-screen.ppm').write_bytes(ppm)\n" - "\n" - "class PillowStyleImage:\n" - " def __init__(self, data):\n" - " self.data = data\n" - " self.size = (1, 1)\n" - "\n" - " def save(self, path):\n" - " pathlib.Path(path).write_bytes(self.data)\n" - "\n" - "def image_open_like_pillow(fp):\n" - " isinstance(fp, pathlib.Path)\n" - " return PillowStyleImage(pathlib.Path(fp).read_bytes())\n" - "\n" - "im = image_open_like_pillow('/tmp/win311-screen.ppm')\n" - "im.save('/tmp/win311-screen.png')\n" - "data_url = 'data:image/png;base64,' + base64.b64encode(\n" - " pathlib.Path('/tmp/win311-screen.png').read_bytes()\n" - ").decode()\n" - "vision = await predict(\n" - " 'image: dspy.Image, question: str -> visible_text: str, answer: str',\n" - " instructions='Inspect this VM screenshot.',\n" - " image=data_url,\n" - " question='Does this show the Windows 3.11 desktop?',\n" - ")\n" - "print(vision.visible_text)\n" - "print(vision.answer)" - ) - }, - ) - - assert tool_call["method"] == "tool_call" - assert tool_call["params"]["name"] == "predict" - assert tool_call["params"]["args"] == [ - "image: dspy.Image, question: str -> visible_text: str, answer: str" - ] - assert tool_call["params"]["kwargs"]["instructions"] == "Inspect this VM screenshot." - assert tool_call["params"]["kwargs"]["question"] == "Does this show the Windows 3.11 desktop?" - image = tool_call["params"]["kwargs"]["image"] - assert image.startswith("data:image/png;base64,") - assert base64.b64decode(image.removeprefix("data:image/png;base64,")) == ( - b"P6\n1 1\n255\n" + bytes([255, 255, 255]) - ) - - assert runner.proc.stdin is not None - assert runner.proc.stdout is not None - runner.proc.stdin.write( - json.dumps( - { - "jsonrpc": "2.0", - "id": tool_call["id"], - "result": { - "type": "json", - "value": '{"visible_text": "desktop", "answer": "yes"}', - }, - } - ) - + "\n" - ) - runner.proc.stdin.flush() - response = json.loads(runner.proc.stdout.readline()) - - assert response["result"]["output"] == "desktop\nyes\n" - - -class TestSbxBackendCreateNaming: - """`sbx create` always receives a known `--name`; the name is never scraped - from stdout (regression guard for issue #39).""" - - def _run_create( - self, tmp_path: Path, *, config: SbxConfig, create_stdout: str - ) -> tuple[SbxBackend, list[str]]: - backend = SbxBackend( - config=config, - preinstall_packages=False, - _supervisor_command=[sys.executable, "-u", str(PAYLOAD_PATH)], - _staging_root=tmp_path / "staging", - ) - captured: dict[str, list[str]] = {} - - def fake_run(cmd, *args, **kwargs): - captured["cmd"] = cmd - return SimpleNamespace(returncode=0, stdout=create_stdout, stderr="") - - with ( - patch("predict_rlm.backends.sbx.backend.shutil.which", return_value="/usr/bin/sbx"), - patch("predict_rlm.backends.sbx.backend.subprocess.run", side_effect=fake_run), - patch.object(SbxBackend, "_prepare_supervisor_script", return_value=tmp_path / "sup.py"), - patch.object(SbxBackend, "_apply_network_policy"), - patch.object(SbxBackend, "_bootstrap_packages"), - patch.object(SbxBackend, "_setup_direct_workspace_aliases_in_sandbox"), - ): - backend._start_sbx_and_prepare_supervisor() - return backend, captured["cmd"] - - def test_generates_name_and_passes_it_to_create(self, tmp_path: Path): - backend, cmd = self._run_create( - tmp_path, config=SbxConfig(), create_stdout="some-auto-name\n" - ) - assert "--name" in cmd - name = cmd[cmd.index("--name") + 1] - assert name.startswith("predict-rlm-") - assert backend._sandbox_name == name - - def test_uses_explicit_config_name(self, tmp_path: Path): - backend, cmd = self._run_create( - tmp_path, config=SbxConfig(name="my-box"), create_stdout="my-box\n" - ) - assert cmd[cmd.index("--name") + 1] == "my-box" - assert backend._sandbox_name == "my-box" - - def test_ignores_update_banner_in_stdout(self, tmp_path: Path): - # The sbx update banner draws a Unicode box; the old code grabbed its - # bottom border as the name. The name must come from `--name`, not stdout. - banner = ( - "╭──────────────────────────────╮\n" - "│ A new version of sbx is out │\n" - "╰──────────────────────────────╯\n" - ) - backend, cmd = self._run_create( - tmp_path, config=SbxConfig(), create_stdout=banner - ) - name = cmd[cmd.index("--name") + 1] - assert name.startswith("predict-rlm-") - assert backend._sandbox_name == name - assert "╰" not in backend._sandbox_name - class TestSbxBackendLocalRunner: def make_interpreter( @@ -987,188 +523,6 @@ def make_interpreter( _staging_root=tmp_path / "staging", ) - def test_execute_and_state_persistence(self, tmp_path: Path): - interpreter = self.make_interpreter(tmp_path) - try: - assert interpreter.execute("x = 7\nprint(x)") == "7\n" - assert interpreter.execute("x += 1\nprint(x)") == "8\n" - finally: - interpreter.shutdown() - - def test_submit_returns_final_output(self, tmp_path: Path): - interpreter = self.make_interpreter(tmp_path) - try: - result = interpreter.execute("SUBMIT(answer='ok')") - finally: - interpreter.shutdown() - - assert isinstance(result, FinalOutput) - assert result.output == {"answer": "ok"} - - def test_debug_logs_runner_and_request_events(self, tmp_path: Path, caplog, capsys): - caplog.set_level(logging.DEBUG, logger="predict_rlm") - interpreter = self.make_interpreter(tmp_path, debug=True) - try: - assert interpreter.execute("print('hi')") == "hi\n" - finally: - interpreter.shutdown() - - stderr = capsys.readouterr().err - assert "output:" not in stderr - events = [record.getMessage().split()[0] for record in caplog.records] - assert "sbx.runner.start" in events - assert "sbx.runner.started" in events - assert "sbx.request.start" in events - assert "sbx.request.ok" in events - assert "sbx.shutdown.complete" in events - - def test_execute_raises_recoverable_code_errors(self, tmp_path: Path): - interpreter = self.make_interpreter(tmp_path) - try: - with pytest.raises(CodeInterpreterError, match="NameError"): - interpreter.execute("print(missing_name)") - finally: - interpreter.shutdown() - - def test_execute_error_includes_partial_output(self, tmp_path: Path): - interpreter = self.make_interpreter(tmp_path) - try: - with pytest.raises(CodeInterpreterError) as exc_info: - interpreter.execute("print('before failure')\nraise ValueError('bad')") - finally: - interpreter.shutdown() - - assert "before failure" in str(exc_info.value) - assert "ValueError" in str(exc_info.value) - assert getattr(exc_info.value, "partial_output") == "before failure\n" - - def test_debug_logs_partial_output_on_error(self, tmp_path: Path, caplog): - caplog.set_level(logging.DEBUG, logger="predict_rlm") - interpreter = self.make_interpreter(tmp_path, debug=True) - try: - with pytest.raises(CodeInterpreterError): - interpreter.execute("print('before failure')\nraise ValueError('bad')") - finally: - interpreter.shutdown() - - messages = "\n".join(record.getMessage() for record in caplog.records) - assert "sandbox.partial_output" in messages - assert "before failure" in messages - - # TEMPORARY SKIP: this exercises SbxBackend over its stdin/stdout transport, which is - # test-only and deprecated (real SBX uses websocket since the exec->ws migration). On - # loaded CI runners the supervisor request intermittently hangs/times out. The verbose - # behavior is unique to this class; re-enable by migrating it to the websocket runner - # when we delete the dead stdin/stdout transport (follow-up). - # local-only: flaky on CI; deprecated SbxBackend stdin/stdout transport, slated for removal - @pytest.mark.local - def test_verbose_prints_output_tool_calls_and_errors(self, tmp_path: Path, capsys): - async def add(a: int, b: int) -> dict: - await asyncio.sleep(0) - return {"total": a + b} - - interpreter = self.make_interpreter( - tmp_path, - verbose=True, - tools={"add": add}, - ) - try: - output = interpreter.execute("result = await add(2, 3)\nprint(result['total'])") - with pytest.raises(CodeInterpreterError, match="ValueError"): - interpreter.execute("raise ValueError('bad')") - finally: - interpreter.shutdown() - - assert output.strip() == "5" - stderr = capsys.readouterr().err - assert "[INFO]" not in stderr - assert "predict_rlm.trace" not in stderr - assert "Tool: add(" in stderr - assert '"args": [2, 3]' in stderr - assert "output:" in stderr - assert "5" in stderr - assert "error (ValueError):" in stderr - assert "bad" in stderr - - def test_verbose_prints_partial_output_before_error(self, tmp_path: Path, capsys): - interpreter = self.make_interpreter(tmp_path, verbose=True) - try: - with pytest.raises(CodeInterpreterError, match="ValueError"): - interpreter.execute("print('before failure')\nraise ValueError('bad')") - finally: - interpreter.shutdown() - - stderr = capsys.readouterr().err - assert "output:" in stderr - assert "before failure" in stderr - assert "error (ValueError):" in stderr - assert "bad" in stderr - - def test_verbose_prints_submit_payload(self, tmp_path: Path, capsys): - interpreter = self.make_interpreter(tmp_path, verbose=True) - try: - result = interpreter.execute("SUBMIT(answer='ok')") - finally: - interpreter.shutdown() - - assert isinstance(result, FinalOutput) - stderr = capsys.readouterr().err - assert "[INFO]" not in stderr - assert "predict_rlm.trace" not in stderr - assert "output:" in stderr - assert '"answer": "ok"' in stderr - - def test_execute_timeout_returns_recoverable_observation( - self, tmp_path: Path - ): - interpreter = self.make_interpreter(tmp_path) - try: - timeout_result = interpreter.execute( - "import sys\n" - "print('before timeout')\n" - "print('stderr before timeout', file=sys.stderr)\n" - "partial_timeout_state = 42\n" - "while True:\n" - " pass\n", - timeout=0.1, - ) - followup = interpreter.execute( - "print('partial_timeout_state' in globals())\nprint('still alive')" - ) - finally: - interpreter.shutdown() - - assert "[Timeout] Iteration execution timed out after 0.1s" in timeout_result - assert "[stdout]\nbefore timeout" in timeout_result - assert "[stderr]\nstderr before timeout" in timeout_result - assert timeout_result.timeout_seconds == 0.1 - assert timeout_result.state == { - "preserved": True, - "source": "live_kernel", - "scope": "full_live", - } - assert timeout_result.state_preserved is True - assert followup == "True\nstill alive\n" - - def test_default_recoverable_timeout_grace_is_shared(self, tmp_path: Path): - from predict_rlm.execution_timeout import ( - DEFAULT_RECOVERABLE_EXECUTION_TIMEOUT_GRACE_SECONDS, - ITERATION_TIMEOUT_FAILURE_CLASS, - ) - - interpreter = self.make_interpreter(tmp_path) - try: - assert DEFAULT_RECOVERABLE_EXECUTION_TIMEOUT_GRACE_SECONDS == 30.0 - assert ( - interpreter._host_watchdog_timeout( - 2.0, - ITERATION_TIMEOUT_FAILURE_CLASS, - ) - == 32.0 - ) - finally: - interpreter.shutdown() - def test_delayed_structured_timeout_uses_recoverable_grace(self, tmp_path: Path): runner_script = tmp_path / "delayed_timeout_runner.py" runner_script.write_text( @@ -1315,65 +669,6 @@ def test_iteration_timeout_recovery_failure_is_bounded_by_grace( assert 0.25 <= time.monotonic() - start < 1.0 - def test_file_helpers_round_trip_virtual_paths(self, tmp_path: Path): - interpreter = self.make_interpreter(tmp_path) - source = tmp_path / "input.txt" - source.write_text("hello", encoding="utf-8") - output = tmp_path / "output.txt" - try: - interpreter.mount_file_at(str(source), "/sandbox/input/source/input.txt") - interpreter.mkdir_p("/sandbox/output/result") - result = interpreter.execute( - "from pathlib import Path\n" - "print(Path(input_path).exists())\n" - "text = Path(input_path).read_text()\n" - "with open('/sandbox/output/result/output.txt', 'w', encoding='utf-8') as f:\n" - " f.write(text + ' sbx')", - variables={"input_path": "/sandbox/input/source/input.txt"}, - ) - files = interpreter.list_dir("/sandbox/output/result") - interpreter.sync_file_to("/sandbox/output/result/output.txt", str(output)) - finally: - interpreter.shutdown() - - assert "True" in result - assert files == ["/sandbox/output/result/output.txt"] - assert output.read_text(encoding="utf-8") == "hello sbx" - - def test_file_helpers_are_host_side_and_do_not_start_runner(self, tmp_path: Path): - interpreter = SbxBackend( - config=SbxConfig(name="file-only-test"), - preinstall_packages=False, - _supervisor_command=[sys.executable, "-c", "raise SystemExit(99)"], - _staging_root=tmp_path / "staging", - ) - source = tmp_path / "host-input.txt" - source.write_text("host visible", encoding="utf-8") - output = tmp_path / "host-output.txt" - - try: - interpreter.mount_file_at(str(source), "/sandbox/input/source/input.txt") - interpreter.mkdir_p("/sandbox/output/result/nested") - staged_output = ( - tmp_path / "staging" / "sandbox" / "output" / "result" / "nested" / "output.txt" - ) - staged_output.write_text("from staging", encoding="utf-8") - - files = interpreter.list_dir("/sandbox/output/result") - interpreter.sync_file_to( - "/sandbox/output/result/nested/output.txt", - str(output), - ) - finally: - interpreter.shutdown() - - staged_input = tmp_path / "staging" / "sandbox" / "input" / "source" / "input.txt" - assert staged_input.read_text(encoding="utf-8") == "host visible" - assert (tmp_path / "staging" / "sandbox" / "output" / "result" / "nested").is_dir() - assert files == ["/sandbox/output/result/nested/output.txt"] - assert output.read_text(encoding="utf-8") == "from staging" - assert interpreter._proc is None - def test_shutdown_removes_owned_staging_root( self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch ): @@ -1416,58 +711,6 @@ def test_shutdown_preserves_caller_owned_staging_root(self, tmp_path: Path): encoding="utf-8" ) == "host visible" - def test_persist_preserves_owned_staging_root( - self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch - ): - monkeypatch.chdir(tmp_path) - interpreter = SbxBackend( - config=SbxConfig(name="local-test", persist=True), - preinstall_packages=False, - _supervisor_command=[sys.executable, "-u", str(PAYLOAD_PATH)], - ) - staging_root = interpreter._staging_root - - try: - interpreter.mkdir_p("/sandbox/output/result") - finally: - interpreter.shutdown() - - assert staging_root.is_dir() - - def test_execute_can_call_registered_host_tools(self, tmp_path: Path): - async def add(a: int, b: int) -> dict: - await asyncio.sleep(0) - return {"total": a + b} - - interpreter = SbxBackend( - config=SbxConfig(name="local-test"), - tools={"add": add}, - preinstall_packages=False, - _supervisor_command=[sys.executable, "-u", str(PAYLOAD_PATH)], - _staging_root=tmp_path / "staging", - ) - try: - output = interpreter.execute("result = await add(2, 3)\nprint(result['total'])") - finally: - interpreter.shutdown() - - assert output.strip() == "5" - - def test_predict_rlm_recovers_after_user_exceptions_and_tools_still_work( - self, tmp_path: Path - ): - pool = SbxPool( - size=1, - config=SbxConfig(name="local-test-user-exceptions"), - preinstall_packages=False, - _supervisor_command=[sys.executable, "-u", str(PAYLOAD_PATH)], - _staging_root=tmp_path / "staging", - ) - try: - assert_predict_rlm_recovers_after_user_exceptions_and_tools_still_work(pool) - finally: - pool.shutdown() - def test_host_tool_synced_file_writeback_updates_sandbox_file(self, tmp_path: Path): received_paths: list[str] = [] @@ -1533,185 +776,6 @@ def mutate(path: Annotated[str, SyncedFile(writeback=False)]) -> str: assert output.strip().splitlines() == ["host only", "sandbox"] - def test_host_tool_synced_file_host_dir_writeback_uses_configured_directory( - self, tmp_path: Path - ): - host_dir = tmp_path / "synced-host-dir" - received_paths: list[str] = [] - - def mutate(path: str) -> str: - received_paths.append(path) - file_path = Path(path) - file_path.write_text( - file_path.read_text(encoding="utf-8") + " + configured", - encoding="utf-8", - ) - return str(file_path) - - mutate.__annotations__["path"] = Annotated[str, SyncedFile(host_dir=str(host_dir))] - - interpreter = SbxBackend( - config=SbxConfig(name="local-test"), - tools={"mutate": mutate}, - preinstall_packages=False, - _supervisor_command=[sys.executable, "-u", str(PAYLOAD_PATH)], - _staging_root=tmp_path / "staging", - ) - source = tmp_path / "input.txt" - source.write_text("sandbox", encoding="utf-8") - try: - interpreter.mount_file_at(str(source), "/sandbox/input/source/input.txt") - output = interpreter.execute( - "received = await mutate(path='/sandbox/input/source/input.txt')\n" - "print(received)\n" - "with open('/sandbox/input/source/input.txt', encoding='utf-8') as f:\n" - " print(f.read())" - ) - finally: - interpreter.shutdown() - - assert received_paths == [str(host_dir / "input.txt")] - assert output.strip().splitlines() == [ - str(host_dir / "input.txt"), - "sandbox + configured", - ] - assert (host_dir / "input.txt").read_text(encoding="utf-8") == ("sandbox + configured") - - def test_execute_serializes_concurrent_requests(self, tmp_path: Path): - runner_script = tmp_path / "detect_concurrent_requests.py" - runner_script.write_text( - """ -import json -import select -import sys -import time - - -def send(message): - sys.stdout.write(json.dumps(message) + "\\n") - sys.stdout.flush() - - -while True: - line = sys.stdin.readline() - if not line: - break - request = json.loads(line) - request_id = request.get("id") - method = request.get("method") - if method == "shutdown": - send({"jsonrpc": "2.0", "result": {"shutdown": True}, "id": request_id}) - break - if method != "execute": - send({"jsonrpc": "2.0", "result": {}, "id": request_id}) - continue - time.sleep(0.2) - readable, _, _ = select.select([sys.stdin], [], [], 0) - if readable: - send({ - "jsonrpc": "2.0", - "error": { - "code": -32000, - "message": "concurrent request detected", - "data": {"type": "RuntimeError", "args": ["concurrent request detected"]}, - }, - "id": request_id, - }) - continue - send({ - "jsonrpc": "2.0", - "result": {"output": request.get("params", {}).get("code", "") + "\\n"}, - "id": request_id, - }) -""".lstrip(), - encoding="utf-8", - ) - interpreter = SbxBackend( - config=SbxConfig(name="local-test", exec_timeout=2), - preinstall_packages=False, - _supervisor_command=[sys.executable, "-u", str(runner_script)], - _staging_root=tmp_path / "staging", - ) - barrier = threading.Barrier(3) - results: list[str] = [] - errors: list[BaseException] = [] - - def execute(code: str) -> None: - barrier.wait() - try: - results.append(interpreter.execute(code).strip()) - except BaseException as exc: - errors.append(exc) - - threads = [ - threading.Thread(target=execute, args=("first",)), - threading.Thread(target=execute, args=("second",)), - ] - try: - for thread in threads: - thread.start() - barrier.wait() - for thread in threads: - thread.join(timeout=3) - finally: - interpreter.shutdown() - - assert [thread.is_alive() for thread in threads] == [False, False] - assert errors == [] - assert sorted(results) == ["first", "second"] - - def test_concurrent_host_tool_calls_do_not_run_serially(self, tmp_path: Path): - async def slow(value: int) -> int: - await asyncio.sleep(0.35) - return value - - interpreter = SbxBackend( - config=SbxConfig(name="local-test", exec_timeout=3), - tools={"slow": slow}, - preinstall_packages=False, - _supervisor_command=[sys.executable, "-u", str(PAYLOAD_PATH)], - _staging_root=tmp_path / "staging", - ) - try: - interpreter.prewarm() - interpreter.execute("pass") - start = time.monotonic() - output = interpreter.execute( - "import asyncio\n" - "results = await asyncio.gather(slow(1), slow(2))\n" - "print(results)" - ) - elapsed = time.monotonic() - start - finally: - interpreter.shutdown() - - assert output.strip() == "[1, 2]" - assert elapsed < 0.6 - - def test_same_interpreter_tool_reentry_raises_runtimeerror(self, tmp_path: Path): - observed_errors: list[str] = [] - - def reenter() -> str: - with pytest.raises(RuntimeError, match="host tool callback") as exc_info: - interpreter.execute("print('nested')") - observed_errors.append(str(exc_info.value)) - return "blocked" - - interpreter = SbxBackend( - config=SbxConfig(name="local-test", exec_timeout=3), - tools={"reenter": reenter}, - preinstall_packages=False, - _supervisor_command=[sys.executable, "-u", str(PAYLOAD_PATH)], - _staging_root=tmp_path / "staging", - ) - try: - output = interpreter.execute("result = await reenter()\nprint(result)") - finally: - interpreter.shutdown() - - assert output.strip() == "blocked" - assert len(observed_errors) == 1 - def test_request_timeout_fires_when_runner_stays_silent(self, tmp_path: Path): interpreter = SbxBackend( config=SbxConfig(name="silent-test", exec_timeout=0.2), @@ -1793,127 +857,7 @@ def make_interpreter( _staging_root=staging_root or tmp_path / "ws-staging", ) - def test_websocket_execute_and_state_persistence(self, tmp_path: Path): - interpreter = self.make_interpreter(tmp_path) - try: - assert interpreter.execute("x = 7\nprint(x)") == "7\n" - assert interpreter.execute("x += 1\nprint(x)") == "8\n" - finally: - interpreter.shutdown() - - def test_pydantic_model_variable_injected_as_plain_data(self, tmp_path: Path): - """A pydantic-model input variable must cross the boundary as plain data. - - Input variables are injected into the sandbox via repr(); a model's repr is a - constructor call (e.g. RfpAnalysis(...)) referencing a class the sandbox does - not have, so it would raise NameError. to_plain_data() normalizes the model to - a dict first. Regression for chained-RLM runs that pass a model output as the - next RLM's input. - """ - from pydantic import BaseModel - - class KeyDate(BaseModel): - name: str - - class RfpAnalysis(BaseModel): - title: str - key_dates: list[KeyDate] - - interpreter = self.make_interpreter(tmp_path) - try: - output = interpreter.execute( - "print(type(rfp).__name__, rfp['title'], rfp['key_dates'][0]['name'])", - variables={"rfp": RfpAnalysis(title="T", key_dates=[KeyDate(name="due")])}, - ) - finally: - interpreter.shutdown() - - assert output == "dict T due\n" - - def test_pydantic_enum_variable_is_plain_data(self, tmp_path: Path): - from enum import Enum - - from pydantic import BaseModel - - class WorkspaceMode(str, Enum): - DIRECT = "direct" - - class PriorTurn(BaseModel): - workspace_mode: WorkspaceMode - - interpreter = self.make_interpreter(tmp_path) - try: - output = interpreter.execute( - "print(prior_turn['workspace_mode'])", - variables={"prior_turn": PriorTurn(workspace_mode=WorkspaceMode.DIRECT)}, - ) - finally: - interpreter.shutdown() - - assert output == "direct\n" - - def test_websocket_host_tool_round_trip(self, tmp_path: Path): - def add(a: int, b: int) -> dict: - return {"total": a + b} - - interpreter = self.make_interpreter(tmp_path, tools={"add": add}) - try: - output = interpreter.execute("result = await add(2, 3)\nprint(result['total'])") - finally: - interpreter.shutdown() - - assert output == "5\n" - - def test_websocket_concurrent_host_tool_timeout_recovers_and_shutdowns( - self, tmp_path: Path - ): - def slow_tool() -> str: - time.sleep(5) - return "slow" - - interpreter = self.make_interpreter(tmp_path, tools={"slow_tool": slow_tool}) - shutdown_duration = None - try: - timeout_result = interpreter.execute( - "import asyncio\n" - "await asyncio.gather(slow_tool(), slow_tool())\n", - timeout=0.1, - ) - output = interpreter.execute("print('still alive')") - shutdown_start = time.perf_counter() - interpreter.shutdown() - shutdown_duration = time.perf_counter() - shutdown_start - finally: - interpreter.shutdown() - - assert "[Timeout] Iteration execution timed out after 0.1s" in timeout_result - assert output == "still alive\n" - assert interpreter._pending_tool_calls == {} - assert shutdown_duration is not None and shutdown_duration < 2 - - def test_websocket_large_host_tool_payload_round_trips(self, tmp_path: Path): - seen_lengths: list[int] = [] - - def predict(signature: str, **kwargs) -> dict: - seen_lengths.append(len(kwargs["text"])) - return {"answer": "4"} - - interpreter = self.make_interpreter(tmp_path, tools={"predict": predict}) - try: - output = interpreter.execute( - "payload = 'x' * 950000\n" - "result = await predict('text: str -> answer: str', text=payload)\n" - "print(result['answer'])" - ) - finally: - interpreter.shutdown() - - assert output == "4\n" - assert seen_lengths == [950000] - - def test_reusable_named_websocket_supervisors_run_concurrently( - self, tmp_path: Path - ): + def test_reusable_named_websocket_supervisors_run_concurrently(self, tmp_path: Path): staging_root = tmp_path / "shared-staging" first = self.make_interpreter( tmp_path, @@ -1937,10 +881,7 @@ def execute(interpreter: SbxBackend, label: str) -> None: try: barrier.wait(timeout=2) outputs[label] = interpreter.execute( - f"owner = {label!r}\n" - "import time\n" - "time.sleep(0.2)\n" - "print(owner)" + f"owner = {label!r}\nimport time\ntime.sleep(0.2)\nprint(owner)" ) except BaseException as exc: errors.append(exc) @@ -1968,127 +909,62 @@ def execute(interpreter: SbxBackend, label: str) -> None: first.shutdown() second.shutdown() - def test_predict_forwards_nested_pydantic_schemas_for_custom_types(self, tmp_path: Path): - """predict() with a custom output type that nests sibling models. - - The host builds the structured-output signature and can't resolve a - REPL-defined type by name, so the sandbox extracts model_json_schema() - and forwards it (on par with JSPI). The model is defined at REPL top - level but predict() is called from inside a function under gather(), and - the type nests sibling models -- which can't resolve via the mismatched - __main__ unless rebuilt against the execution globals. This is the exact - real-world failure mode; without the fix the host falls back to a plain - string signature and the predict() call raises. + def test_predict_result_reconstructs_nested_pydantic_instances(self, tmp_path: Path): + """Custom output types arrive as dicts and are revived to instances. + + The host serializes model instances to dicts for transport. The sandbox + rebuilds them so nested ``item.name`` attribute access works, matching + the JSPI backend and the core instructions for Pydantic return values. """ - received: dict = {} def predict(signature: str, **kwargs) -> dict: - received["schemas"] = kwargs.get("pydantic_schemas") - return {"analysis": {"page_number": 1, "items": [{"name": "x"}]}} + return {"analysis": {"page_number": 2, "items": [{"name": "x"}, {"name": "y"}]}} interpreter = self.make_interpreter(tmp_path, tools={"predict": predict}) try: output = interpreter.execute( - "import asyncio\n" "from pydantic import BaseModel, Field\n" "class PageItem(BaseModel):\n" " name: str\n" "class PageAnalysis(BaseModel):\n" " page_number: int\n" " items: list[PageItem] = Field(default_factory=list)\n" - "async def one(i):\n" - " return await predict('doc: str -> analysis: PageAnalysis', doc='hi')\n" - "await asyncio.gather(*[one(i) for i in range(3)])\n" - "print('ok')" + "r = await predict('doc: str -> analysis: PageAnalysis', doc='hi')\n" + "print(r.analysis.page_number, [i.name for i in r.analysis.items])" ) finally: interpreter.shutdown() - assert output == "ok\n" - schemas = received["schemas"] - assert schemas is not None and "PageAnalysis" in schemas - # nested sibling model must be present in the forwarded schema's $defs - assert "PageItem" in json.dumps(schemas["PageAnalysis"]) + assert output == "2 ['x', 'y']\n" - def test_predict_result_supports_attribute_and_subscript_access(self, tmp_path: Path): - """predict() return must work as ``result.page`` and ``result["page"]``. + def test_predict_reconstruction_preserves_extra_lm_fields(self, tmp_path: Path): + """Deno parity: fields the LM returns beyond the declared model survive. - The host returns a plain dict over the wire, which only supports - subscript -- ``result.page`` would raise ``'dict' object has no - attribute 'page'``. The sandbox wraps it in a Prediction-like object so - both forms work, matching the core instructions and the JSPI backend. + Reconstruction validates into an ``extra='allow'`` subclass (matching the + JSPI/Deno backend) so an unexpected field like ``bonus`` is kept and + attribute-accessible rather than dropped. With a plain (extra='ignore') + model ``r.item.bonus`` would raise AttributeError. """ def predict(signature: str, **kwargs) -> dict: - return {"page": "p1", "items": ["a", "b"]} + return {"item": {"name": "x", "bonus": "kept"}} interpreter = self.make_interpreter(tmp_path, tools={"predict": predict}) try: output = interpreter.execute( - "r = await predict('doc: str -> page: str, items: list[str]', doc='hi')\n" - "print(r.page, r['page'], r.items, r['items'])" + "from pydantic import BaseModel\n" + "class Item(BaseModel):\n" + " name: str\n" + "r = await predict('doc: str -> item: Item', doc='hi')\n" + "print(type(r.item).__name__, r.item.name, r.item.bonus)" ) finally: interpreter.shutdown() - assert output == "p1 p1 ['a', 'b'] ['a', 'b']\n" + assert output == "Item x kept\n" - def test_predict_result_reconstructs_nested_pydantic_instances(self, tmp_path: Path): - """Custom output types arrive as dicts and are revived to instances. - - The host serializes model instances to dicts for transport. The sandbox - rebuilds them so nested ``item.name`` attribute access works, matching - the JSPI backend and the core instructions for Pydantic return values. - """ - - def predict(signature: str, **kwargs) -> dict: - return {"analysis": {"page_number": 2, "items": [{"name": "x"}, {"name": "y"}]}} - - interpreter = self.make_interpreter(tmp_path, tools={"predict": predict}) - try: - output = interpreter.execute( - "from pydantic import BaseModel, Field\n" - "class PageItem(BaseModel):\n" - " name: str\n" - "class PageAnalysis(BaseModel):\n" - " page_number: int\n" - " items: list[PageItem] = Field(default_factory=list)\n" - "r = await predict('doc: str -> analysis: PageAnalysis', doc='hi')\n" - "print(r.analysis.page_number, [i.name for i in r.analysis.items])" - ) - finally: - interpreter.shutdown() - - assert output == "2 ['x', 'y']\n" - - def test_predict_reconstruction_preserves_extra_lm_fields(self, tmp_path: Path): - """Deno parity: fields the LM returns beyond the declared model survive. - - Reconstruction validates into an ``extra='allow'`` subclass (matching the - JSPI/Deno backend) so an unexpected field like ``bonus`` is kept and - attribute-accessible rather than dropped. With a plain (extra='ignore') - model ``r.item.bonus`` would raise AttributeError. - """ - - def predict(signature: str, **kwargs) -> dict: - return {"item": {"name": "x", "bonus": "kept"}} - - interpreter = self.make_interpreter(tmp_path, tools={"predict": predict}) - try: - output = interpreter.execute( - "from pydantic import BaseModel\n" - "class Item(BaseModel):\n" - " name: str\n" - "r = await predict('doc: str -> item: Item', doc='hi')\n" - "print(type(r.item).__name__, r.item.name, r.item.bonus)" - ) - finally: - interpreter.shutdown() - - assert output == "Item x kept\n" - - def test_predict_reconstruction_raises_on_invalid_model_output(self, tmp_path: Path): - """A predict() output the declared model rejects must surface loudly. + def test_predict_reconstruction_raises_on_invalid_model_output(self, tmp_path: Path): + """A predict() output the declared model rejects must surface loudly. When the host returns data that can't satisfy the model (here: missing the required ``name``), reconstruction lets the validation error propagate so the @@ -2117,81 +993,7 @@ def predict(signature: str, **kwargs) -> dict: assert "validation error" in message.lower() assert "'dict' object has no attribute" not in message - def test_predict_reconstruction_gap_function_local_model_not_resolved( - self, tmp_path: Path - ): - """KNOWN GAP: a predict output model defined inside a function isn't revived. - - Reconstruction resolves the output model from the kernel's module globals only - (no call-stack walk). A model defined at REPL top level -- every real example -- - resolves fine, but one defined INSIDE a function is not in module globals, so - the field stays a plain dict and attribute access raises. This is a deliberate, - documented limitation kept out of the hot path: if it ever shows up it fails - loudly here (not silently), and the fix would be stack-aware resolution at the - predict() call site. This test pins the gap so a future change is intentional. - """ - - def predict(signature: str, **kwargs) -> dict: - return {"item": {"name": "x"}} - - interpreter = self.make_interpreter(tmp_path, tools={"predict": predict}) - try: - # Item is function-local -> absent from kernel module globals -> not revived, - # so res.item is a plain dict and res.item.name raises. - with pytest.raises(SandboxExecutionError) as excinfo: - interpreter.execute( - "async def run():\n" - " from pydantic import BaseModel\n" - " class Item(BaseModel):\n" - " name: str\n" - " res = await predict('doc: str -> item: Item', doc='hi')\n" - " return res.item.name\n" - "await run()" - ) - finally: - interpreter.shutdown() - - assert "'dict' object has no attribute 'name'" in str(excinfo.value) - - def test_predict_reconstructs_single_model_under_gather(self, tmp_path: Path): - """Real-world repro: predict() inside a function fanned out via gather(). - - Matches the RFP-page pattern: a single ``-> notes: PageRfpNotes`` output, - the model defined at REPL top level with Optional + list fields, predict() - called inside an async helper, all 38 fanned out with asyncio.gather, then - ``res.notes.page`` accessed. Without nested reconstruction this raises - ``'dict' object has no attribute 'page'``. - """ - - def predict(signature: str, **kwargs) -> dict: - n = kwargs.get("page_number", 0) - return {"notes": {"page": n, "page_type": "cover", "title_or_heading": None, "key_facts": ["a"]}} - - interpreter = self.make_interpreter(tmp_path, tools={"predict": predict}) - try: - output = interpreter.execute( - "import asyncio\n" - "from pydantic import BaseModel, Field\n" - "from typing import Optional\n" - "class PageRfpNotes(BaseModel):\n" - " page: int\n" - " page_type: str = Field(description='x')\n" - " title_or_heading: Optional[str] = None\n" - " key_facts: list[str] = Field(default_factory=list)\n" - "async def analyze(i):\n" - " res = await predict('page_image: dspy.Image, page_number: int, nav_hint: str -> notes: PageRfpNotes', page_number=i+1, nav_hint='h')\n" - " return res.notes\n" - "notes = await asyncio.gather(*[analyze(i) for i in range(38)])\n" - "print(len(notes), notes[0].page, notes[0].page_type, type(notes[0]).__name__)" - ) - finally: - interpreter.shutdown() - - assert output == "38 1 cover PageRfpNotes\n" - - def test_predicts_orphaned_by_gather_failure_do_not_hang_next_execute( - self, tmp_path: Path - ): + def test_predicts_orphaned_by_gather_failure_do_not_hang_next_execute(self, tmp_path: Path): """A gather() that raises early orphans its other predict() calls. Those tasks are left pending on the kernel loop with tool calls already @@ -2199,6 +1001,7 @@ def test_predicts_orphaned_by_gather_failure_do_not_hang_next_execute( host<->kernel protocol and the *next* predict() hangs to the watchdog. The kernel must cancel orphans between executes so the follow-up works. """ + async def predict(signature: str, **kwargs) -> dict: await asyncio.sleep(0.5) return {"a": "ok"} @@ -2226,20 +1029,6 @@ async def predict(signature: str, **kwargs) -> dict: finally: interpreter.shutdown() - def test_websocket_reset_and_shutdown(self, tmp_path: Path): - interpreter = self.make_interpreter(tmp_path) - try: - assert interpreter.execute("x = 7\nprint(x)") == "7\n" - interpreter.reset() - assert interpreter.execute("print('x' in globals())") == "False\n" - proc = interpreter._proc - interpreter.shutdown() - finally: - interpreter.shutdown() - - assert proc is not None - assert proc.poll() is not None - def test_websocket_auth_path_failure_is_reported(self, tmp_path: Path): interpreter = self.make_interpreter( tmp_path, @@ -2254,7 +1043,7 @@ def test_websocket_auth_path_failure_is_reported(self, tmp_path: Path): interpreter.shutdown() @pytest.mark.asyncio - async def _async_operations_do_not_delegate_to_sync_or_to_thread( + async def test_async_operations_do_not_delegate_to_sync_or_to_thread( self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch ): async def add(a: int, b: int) -> dict: @@ -2300,35 +1089,7 @@ def forbidden(*args, **kwargs): assert output.read_text(encoding="utf-8") == "done" @pytest.mark.asyncio - async def _ainterrupt_does_not_call_sync_interrupt_or_to_thread( - self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch - ): - interpreter = self.make_interpreter(tmp_path, startup_timeout=5) - - def forbidden(*args, **kwargs): - raise AssertionError("async interrupt delegated to a sync API") - - monkeypatch.setattr(interpreter, "interrupt", forbidden) - monkeypatch.setattr(asyncio, "to_thread", forbidden) - try: - await interpreter.aexecute("kept = 99") - task = asyncio.create_task( - interpreter.aexecute("import time\ntime.sleep(120)\nprint('done')") - ) - while not interpreter._execution_gate.is_running(): - await asyncio.sleep(0.01) - await asyncio.sleep(0.2) - - assert await interpreter.ainterrupt(timeout=10) is True - result = await task - - assert "done" not in str(result) - assert await interpreter.aexecute("print(kept)") == "99\n" - finally: - await interpreter.ashutdown() - - @pytest.mark.asyncio - async def _aexecute_cancellation_uses_native_interrupt_and_recovers( + async def test_aexecute_cancellation_uses_native_interrupt_and_recovers( self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch ): interpreter = self.make_interpreter(tmp_path, startup_timeout=5) @@ -2357,19 +1118,6 @@ def forbidden(*args, **kwargs): await interpreter.ashutdown() -class TestSbxBackendAsyncNative: - make_interpreter = TestSbxBackendLocalWebSocketRunner.make_interpreter - test_async_operations_do_not_delegate_to_sync_or_to_thread = ( - TestSbxBackendLocalWebSocketRunner._async_operations_do_not_delegate_to_sync_or_to_thread - ) - test_ainterrupt_does_not_call_sync_interrupt_or_to_thread = ( - TestSbxBackendLocalWebSocketRunner._ainterrupt_does_not_call_sync_interrupt_or_to_thread - ) - test_aexecute_cancellation_uses_native_interrupt_and_recovers = ( - TestSbxBackendLocalWebSocketRunner._aexecute_cancellation_uses_native_interrupt_and_recovers - ) - - @pytest.mark.asyncio async def test_cancelled_async_tool_tasks_remain_quarantined_until_done(tmp_path: Path): interpreter = SbxBackend.__new__(SbxBackend) @@ -2495,53 +1243,12 @@ def test_execute_preserves_primary_error_when_post_hook_fails(): assert isinstance(raised.value.post_execute_error, OSError) -def test_owned_sbx_retirement_finishes_before_sync_loop_teardown(monkeypatch): - interpreter = SbxBackend.__new__(SbxBackend) - interpreter._async_pending_tool_calls = {} - interpreter._quarantined_async_tool_calls = set() - interpreter._pending_tool_calls = {} - interpreter._quarantined_tool_calls = set() - cleanup_finished = threading.Event() - shutdown_finished = threading.Event() - - async def stubborn_tool_task(): - try: - await asyncio.Future() - except asyncio.CancelledError: - await asyncio.sleep(0) - cleanup_finished.set() - - async def shutdown(): - shutdown_finished.set() - - interpreter.ashutdown = shutdown # type: ignore[method-assign] - monkeypatch.setattr( - "predict_rlm.backends.sbx.execution.SbxBackend", - lambda **kwargs: interpreter, - ) - backend = SbxExecutionBackend() - - async def run_owned_session(): - context = SimpleNamespace(session=None, ownership=None) - async with backend.start(ExecutionSpec(), context): - task = asyncio.create_task(stubborn_tool_task()) - interpreter._async_pending_tool_calls[task] = 1 - await asyncio.sleep(0) - - asyncio.run(run_owned_session()) - - assert cleanup_finished.is_set() - assert shutdown_finished.is_set() - - @pytest.mark.asyncio @pytest.mark.parametrize( "spec", [ ExecutionSpec( - host_directory_mounts=( - HostDirectoryMount("/host/workspace", "/workspace"), - ) + host_directory_mounts=(HostDirectoryMount("/host/workspace", "/workspace"),) ), ExecutionSpec(allowed_domains=("service.internal",)), ExecutionSpec(extra_read_paths=("/host/input",)), @@ -2622,52 +1329,8 @@ async def prepare(tool, args, kwargs): assert not temporary_root.exists() -class TestSbxBackendInterrupt(TestSbxBackendLocalWebSocketRunner): - """On-demand interrupt and cancellation-safe async execution.""" - - @pytest.mark.local - def test_interrupt_unblocks_long_running_cell(self, tmp_path: Path): - interpreter = self.make_interpreter(tmp_path, startup_timeout=5) - try: - running_flag: dict[str, bool] = {} - - def fire_interrupt(): - time.sleep(1.0) - running_flag["was_running"] = interpreter.interrupt(timeout=10.0) - - thread = threading.Thread(target=fire_interrupt) - thread.start() - start = time.monotonic() - result = interpreter.execute("import time\ntime.sleep(120)\nprint('done')") - elapsed = time.monotonic() - start - thread.join(timeout=5) - - assert elapsed < 30, f"interrupt did not unblock promptly: {elapsed:.1f}s" - assert running_flag.get("was_running") is True - assert "done" not in str(result) - - assert interpreter.execute("print('alive')") == "alive\n" - finally: - interpreter.shutdown() - - @pytest.mark.local - def test_interrupt_preserves_warm_state(self, tmp_path: Path): - interpreter = self.make_interpreter(tmp_path, startup_timeout=5) - try: - assert interpreter.execute("kept = 99\nprint(kept)") == "99\n" - - def fire_interrupt(): - time.sleep(1.0) - interpreter.interrupt(timeout=10.0) - - thread = threading.Thread(target=fire_interrupt) - thread.start() - interpreter.execute("import time\ntime.sleep(120)\nprint('done')") - thread.join(timeout=5) - - assert interpreter.execute("print(kept)") == "99\n" - finally: - interpreter.shutdown() +class TestSbxBackendInterrupt: + make_interpreter = TestSbxBackendLocalWebSocketRunner.make_interpreter @pytest.mark.local def test_interrupt_returns_only_after_cell_releases_gate(self, tmp_path: Path): @@ -2688,9 +1351,9 @@ def run_cell() -> None: was_running = interpreter.interrupt(timeout=10.0) assert was_running is True - assert ( - gate.is_running() is False - ), "interrupt returned before the interrupted cell released the gate" + assert gate.is_running() is False, ( + "interrupt returned before the interrupted cell released the gate" + ) worker.join(timeout=5) assert not worker.is_alive() @@ -2698,1758 +1361,52 @@ def run_cell() -> None: finally: interpreter.shutdown() - @pytest.mark.local - def test_interrupt_returns_false_when_idle(self, tmp_path: Path): - interpreter = self.make_interpreter(tmp_path, startup_timeout=5) - try: - interpreter.execute("print('warm')") - assert interpreter.interrupt(timeout=5.0) is False - finally: - interpreter.shutdown() - - @pytest.mark.local - @pytest.mark.asyncio - async def test_aexecute_cancellation_is_prompt_and_keeps_sandbox_warm( - self, tmp_path: Path - ): - interpreter = self.make_interpreter(tmp_path, startup_timeout=5) - try: - await interpreter.aexecute("seed = 5") - task = asyncio.ensure_future( - interpreter.aexecute("import time\ntime.sleep(120)\nprint('done')") - ) - await asyncio.sleep(1.0) - start = time.monotonic() - task.cancel() - with pytest.raises(asyncio.CancelledError): - await task - elapsed = time.monotonic() - start - assert elapsed < 30, f"cancellation did not unwind promptly: {elapsed:.1f}s" - assert await interpreter.aexecute("print(seed)") == "5\n" - finally: - await interpreter.ashutdown() - - -class TestSupervisorPayloadInterruptMethod: - @pytest.mark.local - def test_interrupt_method_acks_running_false_when_idle(self, tmp_path: Path): - import predict_rlm.backends.supervisor._payload as payload - payload._consume_interrupt_request() - result = asyncio.run( - payload._handle_interrupt_request({"id": 1, "method": "interrupt"}) +@pytest.mark.integration +@pytest.mark.skipif( + not _real_sbx_available(), + reason="real Docker Sandboxes tests require PREDICT_RLM_RUN_SBX_TESTS=1, sbx CLI, and sbx login", +) +class TestSbxBackendRealSbxReattach: + def _list_names(self) -> list[str]: + result = subprocess.run( + ["sbx", "ls"], capture_output=True, text=True, check=False, timeout=15 ) - assert result["result"]["running"] is False - payload._consume_interrupt_request() - - -class TestSbxBackendLocalSupervisorInterrupt(TestSbxBackendLocalWebSocketRunner): - @pytest.mark.local - def test_interrupt_method_trips_interrupt_path_while_running(self, tmp_path: Path): - interpreter = self.make_interpreter(tmp_path, startup_timeout=5) - try: - interpreter.execute("flag = 1") - - def fire_interrupt(): - time.sleep(1.0) - interpreter.interrupt(timeout=10.0) - - thread = threading.Thread(target=fire_interrupt) - thread.start() - start = time.monotonic() - interpreter.execute("import time\ntime.sleep(120)") - elapsed = time.monotonic() - start - thread.join(timeout=5) - assert elapsed < 30 - assert interpreter.execute("print(flag)") == "1\n" - finally: - interpreter.shutdown() + return [line.split()[0] for line in result.stdout.splitlines() if line.split()] + def test_persist_reattach_destroy_lifecycle(self): + name = f"predict-rlm-reattach-{os.getpid()}" + config = SbxConfig(name=name, reuse=True) + marker = f"state-{os.getpid()}" -class TestSbxSupervisorSignalIsolation(TestSbxBackendLocalWebSocketRunner): - @pytest.mark.local - def test_supervisor_runs_in_its_own_process_group(self, tmp_path: Path): - interpreter = self.make_interpreter(tmp_path) + first = SbxBackend(config=config, preinstall_packages=False, debug=True) try: - interpreter.execute("x = 1") - proc = interpreter._proc - assert proc is not None and proc.poll() is None - assert os.getpgid(proc.pid) != os.getpgid(0) - finally: - interpreter.shutdown() - - -class TestSbxCommandConstruction: - def test_default_template_uses_explicit_non_docker_shell_template( - self, monkeypatch, tmp_path: Path - ): - commands: list[list[str]] = [] - - def fake_run(command, **kwargs): - commands.append(command) - return subprocess.CompletedProcess(command, 0, stdout="created-name\n", stderr="") - - monkeypatch.setattr(shutil, "which", lambda name: "/usr/local/bin/sbx") - monkeypatch.setattr(subprocess, "run", fake_run) - interpreter = SbxBackend( - config=SbxConfig(name="created-name"), - preinstall_packages=False, - _staging_root=tmp_path / "staging", - ) - - interpreter._start_sbx_and_build_supervisor_command() - - create_cmd = commands[0] - assert SbxConfig().template == DEFAULT_SBX_TEMPLATE - assert create_cmd[create_cmd.index("--template") + 1] == DEFAULT_SBX_TEMPLATE - - def test_custom_template_overrides_default(self, monkeypatch, tmp_path: Path): - commands: list[list[str]] = [] - - def fake_run(command, **kwargs): - commands.append(command) - return subprocess.CompletedProcess(command, 0, stdout="created-name\n", stderr="") - - monkeypatch.setattr(shutil, "which", lambda name: "/usr/local/bin/sbx") - monkeypatch.setattr(subprocess, "run", fake_run) - interpreter = SbxBackend( - config=SbxConfig( - name="created-name", - template="docker.io/example/custom-template:latest", - ), - preinstall_packages=False, - _staging_root=tmp_path / "staging", - ) - - interpreter._start_sbx_and_build_supervisor_command() - - create_cmd = commands[0] - assert create_cmd[create_cmd.index("--template") + 1] == ( - "docker.io/example/custom-template:latest" - ) - - def test_none_template_omits_template_flag(self, monkeypatch, tmp_path: Path): - commands: list[list[str]] = [] + first.prewarm() + first.execute( + "from pathlib import Path\n" + f"Path('/sandbox/persisted.txt').write_text({marker!r})\n" + "print('wrote')" + ) + first.shutdown() + assert name in self._list_names() - def fake_run(command, **kwargs): - commands.append(command) - return subprocess.CompletedProcess(command, 0, stdout="created-name\n", stderr="") + second = SbxBackend(config=config, preinstall_packages=False) + second.prewarm() + out = second.execute( + "from pathlib import Path\nprint(Path('/sandbox/persisted.txt').read_text())" + ) + assert out.strip() == marker + second.shutdown() + assert name in self._list_names() - monkeypatch.setattr(shutil, "which", lambda name: "/usr/local/bin/sbx") - monkeypatch.setattr(subprocess, "run", fake_run) - interpreter = SbxBackend( - config=SbxConfig(name="created-name", template=None), - preinstall_packages=False, - _staging_root=tmp_path / "staging", - ) - - interpreter._start_sbx_and_build_supervisor_command() - - assert "--template" not in commands[0] - - def test_persist_skips_cleanup_but_is_not_create_flag(self, monkeypatch, tmp_path: Path): - commands: list[list[str]] = [] - - def fake_run(command, **kwargs): - commands.append(command) - return subprocess.CompletedProcess(command, 0, stdout="created-name\n", stderr="") - - monkeypatch.setattr(shutil, "which", lambda name: "/usr/local/bin/sbx") - monkeypatch.setattr(subprocess, "run", fake_run) - interpreter = SbxBackend( - config=SbxConfig(name="created-name", persist=True), - preinstall_packages=False, - _staging_root=tmp_path / "staging", - ) - - command = interpreter._start_sbx_and_build_supervisor_command() - interpreter.shutdown() - - create_cmd = commands[0] - assert create_cmd[:3] == ["sbx", "create", "shell"] - assert "--persist" not in create_cmd - assert command[:4] == ["sbx", "exec", "-i", "-w"] - assert not any(cmd[:2] == ["sbx", "rm"] for cmd in commands) - - def test_websocket_supervisor_starts_foreground_sentinel_and_publishes_port( - self, monkeypatch, tmp_path: Path - ): - run_commands: list[list[str]] = [] - popen_commands: list[list[str]] = [] - - class FakeProcess: - stdout = None - stderr = None - stdin = None - pid = 12345 - - def poll(self): - return None - - def fake_run(command, **kwargs): - run_commands.append(command) - return subprocess.CompletedProcess(command, 0, stdout="created-name\n", stderr="") - - def fake_popen(command, **kwargs): - popen_commands.append(command) - return FakeProcess() - - monkeypatch.setattr(shutil, "which", lambda name: "/usr/local/bin/sbx") - monkeypatch.setattr(subprocess, "run", fake_run) - monkeypatch.setattr(subprocess, "Popen", fake_popen) - interpreter = SbxBackend( - config=SbxConfig( - name="created-name", - websocket_port=8766, - websocket_max_message_bytes=32 * 1024 * 1024, - ), - preinstall_packages=False, - _staging_root=tmp_path / "staging", - ) - monkeypatch.setattr( - interpreter, - "_publish_websocket_port", - lambda port=None: "ws://127.0.0.1:49152/test", - ) - - interpreter._start_sbx_websocket_supervisor() - - assert len(popen_commands) == 1 - supervisor_exec = popen_commands[0] - assert supervisor_exec[:4] == ["sbx", "exec", "-w", str(tmp_path / "staging")] - assert "-d" not in supervisor_exec - assert "-i" not in supervisor_exec - assert "--websocket-host" in supervisor_exec - assert supervisor_exec[supervisor_exec.index("--websocket-port") + 1] == "8766" - assert supervisor_exec[supervisor_exec.index("--websocket-max-message-bytes") + 1] == str( - 32 * 1024 * 1024 - ) - assert interpreter._proc is not None - assert not any(cmd[:3] == ["sbx", "exec", "-d"] for cmd in run_commands) - - def test_websocket_supervisor_uses_dynamic_port_by_default( - self, monkeypatch, tmp_path: Path - ): - popen_commands: list[list[str]] = [] - published_ports: list[int | None] = [] - - class FakeProcess: - stdout = None - stderr = None - stdin = None - pid = 12345 - - def poll(self): - return None - - def fake_run(command, **kwargs): - return subprocess.CompletedProcess(command, 0, stdout="created-name\n", stderr="") - - def fake_popen(command, **kwargs): - popen_commands.append(command) - return FakeProcess() - - monkeypatch.setattr(shutil, "which", lambda name: "/usr/local/bin/sbx") - monkeypatch.setattr(subprocess, "run", fake_run) - monkeypatch.setattr(subprocess, "Popen", fake_popen) - monkeypatch.setattr(secrets, "randbelow", lambda upper: 12345) - interpreter = SbxBackend( - config=SbxConfig(name="created-name"), - preinstall_packages=False, - _staging_root=tmp_path / "staging", - ) - - def fake_publish(port=None): - published_ports.append(port) - return "ws://127.0.0.1:49152/test" - - monkeypatch.setattr(interpreter, "_publish_websocket_port", fake_publish) - - interpreter._start_sbx_websocket_supervisor() - - supervisor_exec = popen_commands[0] - assert supervisor_exec[supervisor_exec.index("--websocket-port") + 1] == "32345" - assert published_ports == [32345] - assert interpreter._active_websocket_port == 32345 - - def test_websocket_recovery_restarts_detached_supervisor_after_kill( - self, monkeypatch, tmp_path: Path - ): - commands: list[list[str]] = [] - - def fake_run(command, **kwargs): - commands.append(command) - return subprocess.CompletedProcess(command, 0, stdout="", stderr="") - - monkeypatch.setattr(subprocess, "run", fake_run) - interpreter = SbxBackend( - config=SbxConfig(name="created-name"), - preinstall_packages=False, - _staging_root=tmp_path / "staging", - ) - interpreter._sandbox_name = "created-name" - interpreter._prepared_supervisor_path = tmp_path / "staging" / ".predict_rlm_supervisor" / "_payload.py" - interpreter._websocket_url = "ws://127.0.0.1:49152/predict-rlm/old" - interpreter._published_websocket_url = interpreter._websocket_url - - started: list[bool] = [] - connected: list[str] = [] - - def fake_start_sbx_websocket_supervisor(): - started.append(True) - interpreter._websocket_url = "ws://127.0.0.1:49153/predict-rlm/new" - - def fake_connect_websocket_supervisor(url: str): - connected.append(url) - interpreter._ws = object() - - monkeypatch.setattr( - interpreter, - "_start_sbx_websocket_supervisor", - fake_start_sbx_websocket_supervisor, - ) - monkeypatch.setattr( - interpreter, - "_connect_websocket_supervisor", - fake_connect_websocket_supervisor, - ) - - interpreter._kill_websocket_supervisor() - interpreter._ensure_websocket_supervisor() - - assert interpreter._published_websocket_url is None - assert started == [True] - assert connected == ["ws://127.0.0.1:49153/predict-rlm/new"] - assert any( - cmd[:5] == ["sbx", "exec", "-w", str(tmp_path / "staging"), "created-name"] - for cmd in commands - ) - - def test_published_websocket_endpoint_parses_localhost_port(self, tmp_path: Path): - interpreter = SbxBackend( - preinstall_packages=False, - _staging_root=tmp_path / "staging", - ) - interpreter._websocket_path = "/predict-rlm/token" - - assert ( - interpreter._parse_published_websocket_endpoint( - "Published 8765/tcp to localhost:49152\n" - ) - == "ws://localhost:49152/predict-rlm/token" - ) - assert ( - interpreter._parse_published_websocket_endpoint("http://127.0.0.1:49153") - == "ws://127.0.0.1:49153/predict-rlm/token" - ) - - def test_published_websocket_endpoint_parse_failure_is_fatal(self, tmp_path: Path): - interpreter = SbxBackend( - preinstall_packages=False, - _staging_root=tmp_path / "staging", - ) - - with pytest.raises(SandboxFatalError, match="published WebSocket endpoint"): - interpreter._parse_published_websocket_endpoint("no ports here") - - def test_shutdown_forces_sbx_removal_without_confirmation(self, monkeypatch, tmp_path: Path): - commands: list[list[str]] = [] - - def fake_run(command, **kwargs): - commands.append(command) - return subprocess.CompletedProcess(command, 0, stdout="", stderr="") - - monkeypatch.setattr(subprocess, "run", fake_run) - interpreter = SbxBackend( - config=SbxConfig(name="created-name"), - preinstall_packages=False, - _staging_root=tmp_path / "staging", - ) - interpreter._sandbox_name = "created-name" - - interpreter.shutdown() - - assert ["sbx", "rm", "--force", "created-name"] in commands - - def test_workspace_flags_include_read_only_primary_and_extra_workspaces( - self, monkeypatch, tmp_path: Path - ): - commands: list[list[str]] = [] - extra_one = tmp_path / "extra-one" - extra_two = tmp_path / "extra-two" - - def fake_run(command, **kwargs): - commands.append(command) - return subprocess.CompletedProcess(command, 0, stdout="created-name\n", stderr="") - - monkeypatch.setattr(shutil, "which", lambda name: "/usr/local/bin/sbx") - monkeypatch.setattr(subprocess, "run", fake_run) - interpreter = SbxBackend( - config=SbxConfig( - name="created-name", - workspace_read_only=True, - extra_workspaces=[str(extra_one), str(extra_two)], - ), - preinstall_packages=False, - _staging_root=tmp_path / "staging", - ) - - interpreter._start_sbx_and_build_supervisor_command() - - create_cmd = commands[0] - workspace_arg = f"{tmp_path / 'staging'}:ro" - assert create_cmd[:4] == ["sbx", "create", "shell", workspace_arg] - assert create_cmd[4:6] == [str(extra_one), str(extra_two)] - - def test_direct_workspace_flags_enforce_per_mount_read_only_access( - self, monkeypatch, tmp_path: Path - ): - commands: list[list[str]] = [] - read_only = tmp_path / "dataset" - writable = tmp_path / "repository" - read_only.mkdir() - writable.mkdir() - - def fake_run(command, **kwargs): - commands.append(command) - return subprocess.CompletedProcess( - command, 0, stdout="created-name\n", stderr="" - ) - - monkeypatch.setattr(shutil, "which", lambda name: "/usr/local/bin/sbx") - monkeypatch.setattr(subprocess, "run", fake_run) - interpreter = SbxBackend( - config=SbxConfig(name="created-name"), - preinstall_packages=False, - _staging_root=tmp_path / "staging", - direct_workspace_mounts=[ - DirectWorkspaceMount( - str(read_only), "/datasets/input", read_only=True - ), - DirectWorkspaceMount(str(writable), "/repository"), - ], - ) - - interpreter._start_sbx_and_build_supervisor_command() - - create_cmd = commands[0] - assert f"{read_only}:ro" in create_cmd - assert str(writable) in create_cmd - - def test_running_interpreter_accepts_semantically_reordered_direct_mounts( - self, - monkeypatch, - tmp_path: Path, - ): - first = HostDirectoryMount(str(tmp_path / "first"), "/first") - second = HostDirectoryMount( - str(tmp_path / "second"), - "/second", - read_only=True, - ) - interpreter = SbxBackend( - preinstall_packages=False, - _staging_root=tmp_path / "staging", - direct_workspace_mounts=[first, second], - ) - monkeypatch.setattr(interpreter, "_transport_running", lambda: True) - - interpreter.configure_direct_workspace_mounts([second, first]) - - assert interpreter._direct_workspace_mounts == [first, second] - - def test_default_workspace_is_staging_root_not_repo(self, monkeypatch, tmp_path: Path): - commands: list[list[str]] = [] - - def fake_run(command, **kwargs): - commands.append(command) - return subprocess.CompletedProcess(command, 0, stdout="created-name\n", stderr="") - - monkeypatch.setattr(shutil, "which", lambda name: "/usr/local/bin/sbx") - monkeypatch.setattr(subprocess, "run", fake_run) - interpreter = SbxBackend( - config=SbxConfig(name="created-name"), - preinstall_packages=False, - _staging_root=tmp_path / "staging", - ) - - command = interpreter._start_sbx_and_build_supervisor_command() - - create_cmd = commands[0] - assert create_cmd[:4] == ["sbx", "create", "shell", str(tmp_path / "staging")] - assert str(Path.cwd()) not in create_cmd - assert command[:5] == ["sbx", "exec", "-i", "-w", str(tmp_path / "staging")] - - def test_supervisor_command_uses_python3_executable(self, monkeypatch, tmp_path: Path): - def fake_run(command, **kwargs): - return subprocess.CompletedProcess(command, 0, stdout="created-name\n", stderr="") - - monkeypatch.setattr(shutil, "which", lambda name: "/usr/local/bin/sbx") - monkeypatch.setattr(subprocess, "run", fake_run) - interpreter = SbxBackend( - config=SbxConfig(name="created-name"), - preinstall_packages=False, - _staging_root=tmp_path / "staging", - ) - - command = interpreter._start_sbx_and_build_supervisor_command() - - assert "python" not in command - supervisor_path = tmp_path / "staging" / ".predict_rlm_supervisor" / "_payload.py" - assert supervisor_path.read_text(encoding="utf-8") == PAYLOAD_PATH.read_text( - encoding="utf-8" - ) - assert command[-3:] == ["python3", "-u", str(supervisor_path)] - - def test_runner_restart_reuses_existing_sandbox( - self, monkeypatch, tmp_path: Path - ): - commands: list[list[str]] = [] - - def fake_run(command, **kwargs): - commands.append(command) - return subprocess.CompletedProcess(command, 0, stdout="created-name\n", stderr="") - - monkeypatch.setattr(shutil, "which", lambda name: "/usr/local/bin/sbx") - monkeypatch.setattr(subprocess, "run", fake_run) - interpreter = SbxBackend( - config=SbxConfig(name="created-name"), - preinstall_packages=False, - _staging_root=tmp_path / "staging", - ) - interpreter._sandbox_name = "created-name" - - command = interpreter._start_sbx_and_build_supervisor_command() - - assert commands == [] - assert command[:5] == ["sbx", "exec", "-i", "-w", str(tmp_path / "staging")] - assert "created-name" in command - - def test_package_bootstrap_failure_raises_context(self, monkeypatch, tmp_path: Path): - def fake_run(command, **kwargs): - return subprocess.CompletedProcess( - command, - 17, - stdout="download started", - stderr="no matching distribution", - ) - - monkeypatch.setattr(subprocess, "run", fake_run) - interpreter = SbxBackend( - config=SbxConfig(exec_timeout=1), - skill_packages=["missing-package"], - preinstall_packages=False, - _staging_root=tmp_path / "staging", - ) - interpreter._sandbox_name = "created-name" - - with pytest.raises(SandboxFatalError, match="missing-package"): - interpreter._bootstrap_packages() - - def test_package_bootstrap_uses_docker_sandbox_safe_pip(self, monkeypatch, tmp_path: Path): - commands = [] - - def fake_run(command, **kwargs): - commands.append(command) - return subprocess.CompletedProcess(command, 0, stdout="", stderr="") - - monkeypatch.setattr(subprocess, "run", fake_run) - interpreter = SbxBackend( - preinstall_packages=True, - _staging_root=tmp_path / "staging", - ) - interpreter._sandbox_name = "created-name" - - interpreter._bootstrap_packages() - - assert commands == [ - [ - "sbx", - "exec", - "-w", - str(tmp_path / "staging"), - "created-name", - "python3", - "-m", - "pip", - "install", - "--break-system-packages", - "websockets", - "pydantic", - "pandas", - ] - ] - - -class TestSbxPool: - def test_start_prewarms_interpreters_concurrently(self, tmp_path: Path, monkeypatch): - pool = SbxPool( - size=2, - config=SbxConfig(name="pool-test"), - preinstall_packages=False, - _staging_root=tmp_path / "pool", - ) - barrier = threading.Barrier(2) - active = 0 - max_active = 0 - active_lock = threading.Lock() - prewarmed_indexes: list[int] = [] - - class FakeInterpreter: - def __init__(self, index: int) -> None: - self.index = index - self.shutdown_called = False - - def prewarm(self) -> None: - nonlocal active, max_active - with active_lock: - active += 1 - max_active = max(max_active, active) - try: - barrier.wait(timeout=1) - prewarmed_indexes.append(self.index) - finally: - with active_lock: - active -= 1 - - def shutdown(self) -> None: - self.shutdown_called = True - - monkeypatch.setattr( - pool, - "_create_interpreter", - lambda index: FakeInterpreter(index), - ) - - try: - pool.start() - - assert max_active == 2 - assert prewarmed_indexes == [1, 0] or prewarmed_indexes == [0, 1] - assert [interpreter.index for interpreter in pool._all_interpreters] == [0, 1] - assert pool._started - assert pool._available.qsize() == 2 - finally: - pool.shutdown() - - def test_lease_logging_overrides_do_not_reconfigure_pool( - self, tmp_path: Path, monkeypatch - ): - pool = SbxPool( - size=2, - config=SbxConfig(name="pool-test"), - preinstall_packages=False, - _staging_root=tmp_path / "pool", - ) - created = [] - - class FakeInterpreter: - def __init__(self, index: int) -> None: - self.index = index - self.configure_debug_calls: list[bool] = [] - self.configure_verbose_calls: list[bool] = [] - self.runtime_calls: list[dict] = [] - - def prewarm(self) -> None: - return None - - def configure_debug(self, enabled: bool) -> None: - self.configure_debug_calls.append(enabled) - - def configure_verbose(self, enabled: bool) -> None: - self.configure_verbose_calls.append(enabled) - - def configure_runtime(self, **kwargs) -> None: - self.runtime_calls.append(kwargs) - - def reset(self) -> None: - return None - - def shutdown(self) -> None: - return None - - def create_interpreter(index: int) -> FakeInterpreter: - interpreter = FakeInterpreter(index) - created.append(interpreter) - return interpreter - - monkeypatch.setattr(pool, "_create_interpreter", create_interpreter) - - try: - pool.start() - - with pool.lease(debug=True, verbose=True) as interpreter: - assert interpreter is created[0] - - assert created[0].runtime_calls[-1]["debug"] is True - assert created[0].runtime_calls[-1]["verbose"] is True - assert created[0].configure_debug_calls == [] - assert created[0].configure_verbose_calls == [] - assert created[1].configure_debug_calls == [] - assert created[1].configure_verbose_calls == [] - assert pool.debug is False - assert pool.verbose is False - assert pool._interpreter_kwargs["debug"] is False - assert pool._interpreter_kwargs["verbose"] is False - - with pool.lease() as interpreter: - assert interpreter is created[1] - with pool.lease() as interpreter: - assert interpreter is created[0] - - assert created[0].runtime_calls[-1]["debug"] is False - assert created[0].runtime_calls[-1]["verbose"] is False - finally: - pool.shutdown() - - def test_start_failure_shuts_down_created_interpreters_and_leaves_pool_stopped( - self, tmp_path: Path, monkeypatch - ): - pool = SbxPool( - size=3, - config=SbxConfig(name="pool-test"), - preinstall_packages=False, - _staging_root=tmp_path / "pool", - ) - created = [] - - class FakeInterpreter: - def __init__(self, index: int) -> None: - self.index = index - self.shutdown_called = False - - def prewarm(self) -> None: - if self.index == 1: - raise RuntimeError("prewarm failed") - - def shutdown(self) -> None: - self.shutdown_called = True - - def create_interpreter(index: int) -> FakeInterpreter: - interpreter = FakeInterpreter(index) - created.append(interpreter) - return interpreter - - monkeypatch.setattr(pool, "_create_interpreter", create_interpreter) - - with pytest.raises(RuntimeError, match="prewarm failed"): - pool.start() - - assert created - assert all(interpreter.shutdown_called for interpreter in created) - assert not pool._started - assert pool._all_interpreters == [] - assert pool._available.qsize() == 0 - - def test_shutdown_runs_concurrently_and_attempts_all_interpreters( - self, tmp_path: Path, monkeypatch - ): - pool = SbxPool( - size=3, - config=SbxConfig(name="pool-test"), - preinstall_packages=False, - _staging_root=tmp_path / "pool", - ) - barrier = threading.Barrier(3) - active = 0 - max_active = 0 - active_lock = threading.Lock() - shutdown_indexes: list[int] = [] - - class FakeInterpreter: - def __init__(self, index: int) -> None: - self.index = index - - def prewarm(self) -> None: - return None - - def shutdown(self) -> None: - nonlocal active, max_active - with active_lock: - active += 1 - max_active = max(max_active, active) - try: - barrier.wait(timeout=1) - shutdown_indexes.append(self.index) - if self.index == 1: - raise RuntimeError("shutdown failed") - finally: - with active_lock: - active -= 1 - - monkeypatch.setattr(pool, "_create_interpreter", lambda index: FakeInterpreter(index)) - pool.start() - - with pytest.raises(RuntimeError, match="shutdown failed"): - pool.shutdown() - - assert max_active == 3 - assert sorted(shutdown_indexes) == [0, 1, 2] - assert not pool._started - assert pool._shutdown - assert pool._all_interpreters == [] - assert pool._available.qsize() == 0 - - def test_pool_prewarms_all_interpreters(self, tmp_path: Path): - pool = SbxPool( - size=2, - config=SbxConfig(name="pool-test"), - preinstall_packages=False, - _supervisor_command=[sys.executable, "-u", str(PAYLOAD_PATH)], - _staging_root=tmp_path / "pool", - ) - - try: - pool.start() - - assert len(pool._all_interpreters) == 2 - assert all(interpreter._proc is not None for interpreter in pool._all_interpreters) - assert [interpreter.config.name for interpreter in pool._all_interpreters] == [ - "pool-test-0", - "pool-test-1", - ] - finally: - pool.shutdown() - - def test_pool_assigns_unique_names_when_config_name_is_omitted(self, tmp_path: Path): - pool = SbxPool( - size=2, - preinstall_packages=False, - _supervisor_command=[sys.executable, "-u", str(PAYLOAD_PATH)], - _staging_root=tmp_path / "pool", - ) - - try: - pool.start() - - names = [interpreter.config.name for interpreter in pool._all_interpreters] - assert names == [ - f"{pool._pool_name_prefix}-0", - f"{pool._pool_name_prefix}-1", - ] - assert len(set(names)) == 2 - finally: - pool.shutdown() - - def test_lease_is_exclusive_and_release_resets(self, tmp_path: Path): - pool = SbxPool( - size=1, - config=SbxConfig(name="pool-test"), - preinstall_packages=False, - _supervisor_command=[sys.executable, "-u", str(PAYLOAD_PATH)], - _staging_root=tmp_path / "pool", - ) - acquired = threading.Event() - released = threading.Event() - - def second_lease() -> None: - with pool.lease() as interpreter: - acquired.set() - assert interpreter.execute("print('x' in globals())").strip() == "False" - - try: - pool.start() - with pool.lease() as interpreter: - interpreter.execute("x = 7") - staged = pool._all_interpreters[0]._host_path_for_virtual_path( - "/sandbox/output/value.txt" - ) - staged.parent.mkdir(parents=True, exist_ok=True) - staged.write_text("leaked", encoding="utf-8") - thread = threading.Thread(target=second_lease) - thread.start() - assert not acquired.wait(0.2) - - released.set() - thread.join(timeout=5) - - assert released.is_set() - assert acquired.is_set() - with pool.lease() as interpreter: - assert interpreter.list_dir("/sandbox") == [] - finally: - pool.shutdown() - - def test_shutdown_unblocks_waiting_lease_and_does_not_return_interpreter( - self, tmp_path: Path, monkeypatch - ): - pool = SbxPool( - size=1, - config=SbxConfig(name="pool-test"), - preinstall_packages=False, - _staging_root=tmp_path / "pool", - ) - - class FakeInterpreter: - def __init__(self) -> None: - self.reset_called = False - self.shutdown_called = False - - def prewarm(self) -> None: - return None - - def configure_runtime(self, **kwargs) -> None: - return None - - def reset(self) -> None: - self.reset_called = True - - def shutdown(self) -> None: - self.shutdown_called = True - - interpreter = FakeInterpreter() - monkeypatch.setattr(pool, "_create_interpreter", lambda index: interpreter) - - errors: list[str] = [] - - def waiting_lease() -> None: - try: - with pool.lease(): - errors.append("acquired") - except RuntimeError as exc: - errors.append(str(exc)) - - pool.start() - with pool.lease(): - thread = threading.Thread(target=waiting_lease) - thread.start() - time.sleep(0.1) - - pool.shutdown() - thread.join(timeout=2) - - assert not thread.is_alive() - assert errors == ["SbxPool is shut down"] - assert pool._available.qsize() == 0 - assert pool._all_interpreters == [] - - assert interpreter.shutdown_called - assert not interpreter.reset_called - assert pool._available.qsize() == 0 - - def test_shutdown_requested_during_start_prevents_waiting_lease_acquire( - self, tmp_path: Path, monkeypatch - ): - pool = SbxPool( - size=1, - config=SbxConfig(name="pool-test"), - preinstall_packages=False, - _staging_root=tmp_path / "pool", - ) - prewarm_started = threading.Event() - allow_prewarm = threading.Event() - - class FakeInterpreter: - def __init__(self) -> None: - self.shutdown_called = False - - def prewarm(self) -> None: - prewarm_started.set() - assert allow_prewarm.wait(timeout=2) - - def configure_runtime(self, **kwargs) -> None: - return None - - def reset(self) -> None: - return None - - def shutdown(self) -> None: - self.shutdown_called = True - - interpreter = FakeInterpreter() - monkeypatch.setattr(pool, "_create_interpreter", lambda index: interpreter) - - lease_results: list[str] = [] - - def lease_during_start() -> None: - try: - with pool.lease(): - lease_results.append("acquired") - except RuntimeError as exc: - lease_results.append(str(exc)) - - lease_thread = threading.Thread(target=lease_during_start) - lease_thread.start() - assert prewarm_started.wait(timeout=2) - - shutdown_thread = threading.Thread(target=pool.shutdown) - shutdown_thread.start() - deadline = time.monotonic() + 2 - while time.monotonic() < deadline: - with pool._state_changed: - if pool._shutdown_requested: - break - time.sleep(0.01) - else: - pytest.fail("shutdown did not request pool stop while startup was active") - allow_prewarm.set() - - lease_thread.join(timeout=2) - shutdown_thread.join(timeout=2) - - assert not lease_thread.is_alive() - assert not shutdown_thread.is_alive() - assert lease_results == ["SbxPool is shut down"] - assert interpreter.shutdown_called - assert pool._available.qsize() == 0 - - def test_shutdown_before_lease_autostart_prevents_acquire( - self, tmp_path: Path, monkeypatch - ): - pool = SbxPool( - size=1, - config=SbxConfig(name="pool-test"), - preinstall_packages=False, - _staging_root=tmp_path / "pool", - ) - - class FakeInterpreter: - def prewarm(self) -> None: - return None - - def configure_runtime(self, **kwargs) -> None: - return None - - def reset(self) -> None: - return None - - def shutdown(self) -> None: - return None - - monkeypatch.setattr(pool, "_create_interpreter", lambda index: FakeInterpreter()) - original_begin_start = pool._begin_start - begin_start_entered = threading.Event() - allow_begin_start = threading.Event() - lease_results: list[str] = [] - - def delayed_begin_start(*, allow_restart: bool) -> bool: - assert not allow_restart - begin_start_entered.set() - assert allow_begin_start.wait(timeout=2) - return original_begin_start(allow_restart=allow_restart) - - monkeypatch.setattr(pool, "_begin_start", delayed_begin_start) - - def lease_during_shutdown() -> None: - try: - with pool.lease(): - lease_results.append("acquired") - except RuntimeError as exc: - lease_results.append(str(exc)) - - lease_thread = threading.Thread(target=lease_during_shutdown) - lease_thread.start() - assert begin_start_entered.wait(timeout=2) - - pool.shutdown() - allow_begin_start.set() - lease_thread.join(timeout=2) - - assert not lease_thread.is_alive() - assert lease_results == ["SbxPool is shut down"] - assert pool._available.qsize() == 0 - assert pool._all_interpreters == [] - - def test_lease_after_shutdown_raises_until_explicit_restart( - self, tmp_path: Path, monkeypatch - ): - pool = SbxPool( - size=1, - config=SbxConfig(name="pool-test"), - preinstall_packages=False, - _staging_root=tmp_path / "pool", - ) - - class FakeInterpreter: - def __init__(self, index: int) -> None: - self.index = index - - def prewarm(self) -> None: - return None - - def configure_runtime(self, **kwargs) -> None: - return None - - def reset(self) -> None: - return None - - def shutdown(self) -> None: - return None - - created: list[FakeInterpreter] = [] - - def create_interpreter(index: int) -> FakeInterpreter: - interpreter = FakeInterpreter(index) - created.append(interpreter) - return interpreter - - monkeypatch.setattr(pool, "_create_interpreter", create_interpreter) - - pool.start() - pool.shutdown() - - with pytest.raises(RuntimeError, match="SbxPool is shut down"): - with pool.lease(): - pass - - pool.start() - try: - with pool.lease() as interpreter: - assert interpreter is created[-1] - finally: - pool.shutdown() - - -@pytest.mark.sbx -@pytest.mark.integration -@pytest.mark.skipif( - not _real_sbx_available(), - reason="real Docker Sandboxes tests require PREDICT_RLM_RUN_SBX_TESTS=1, sbx CLI, and sbx login", -) -class TestSbxBackendRealSbx: - def test_real_sbx_executes_basic_python(self): - interpreter = SbxBackend( - config=SbxConfig(name=f"predict-rlm-test-{os.getpid()}"), - preinstall_packages=False, - ) - try: - output = interpreter.execute("print(2 + 3)") - finally: - interpreter.shutdown() - - assert output.strip() == "5" - - @pytest.mark.asyncio - async def test_real_sbx_async_transport_does_not_delegate_to_sync( - self, monkeypatch: pytest.MonkeyPatch - ): - interpreter = SbxBackend( - config=SbxConfig(name=f"predict-rlm-test-async-{os.getpid()}"), - preinstall_packages=False, - ) - - def forbidden(*args, **kwargs): - raise AssertionError("real async SBX delegated to a sync API") - - monkeypatch.setattr(interpreter, "execute", forbidden) - monkeypatch.setattr(interpreter, "interrupt", forbidden) - monkeypatch.setattr(interpreter, "shutdown", forbidden) - monkeypatch.setattr(asyncio, "to_thread", forbidden) - try: - output = await interpreter.aexecute("print(2 + 3)") - finally: - await interpreter.ashutdown() - - assert output.strip() == "5" - - def test_real_sbx_predict_reconstructs_pydantic_output_under_gather(self): - """End-to-end repro of the RFP-page failure in a real sbx sandbox. - - The host serializes a custom output model to a dict for transport; the - sandbox must revive it to a real instance so ``res.insight.page`` works. - This runs reconstruction inside the actual sandbox (its own pip-installed - pydantic) under asyncio.gather -- the exact path that was returning a bare - dict and raising ``'dict' object has no attribute 'page'``. preinstall is - required so pydantic exists in the sandbox. - """ - - def predict(signature: str, **kwargs) -> dict: - n = kwargs.get("page_number", 1) - return { - "insight": { - "page": n, - "title_or_section": "T", - "purpose": "p", - "key_facts": ["a"], - "proposal_requirements": [], - } - } - - interpreter = SbxBackend( - config=SbxConfig(name=f"predict-rlm-test-predict-{os.getpid()}"), - tools={"predict": predict}, - preinstall_packages=True, - ) - try: - output = interpreter.execute( - "import asyncio\n" - "from pydantic import BaseModel\n" - "from typing import Optional\n" - "class PageInsight(BaseModel):\n" - " page: int\n" - " title_or_section: Optional[str] = None\n" - " purpose: str\n" - " key_facts: list[str] = []\n" - " proposal_requirements: list[str] = []\n" - "async def inspect_page(i):\n" - " res = await predict('page: dspy.Image, page_number: int -> insight: PageInsight', page_number=i+1)\n" - " return res.insight\n" - "results = await asyncio.gather(*[inspect_page(i) for i in range(9)])\n" - # Nested reconstructed values are REAL Pydantic instances: attribute - # access (not subscript), isinstance, and model_dump all work. - "print(len(results), type(results[0]).__name__, results[0].page," - " isinstance(results[0], PageInsight), results[0].model_dump()['purpose'])", - timeout=30, - ) - finally: - interpreter.shutdown() - - assert output.strip() == "9 PageInsight 1 True p" - - def test_real_sbx_timeout_is_recoverable_and_runner_survives(self): - interpreter = SbxBackend( - config=SbxConfig(name=f"predict-rlm-test-timeout-{os.getpid()}"), - preinstall_packages=False, - ) - try: - timeout_result = interpreter.execute( - "import sys\n" - "print('before timeout')\n" - "print('stderr before timeout', file=sys.stderr)\n" - "partial_timeout_state = 41\n" - "while True:\n" - " pass\n", - timeout=0.2, - ) - followup = interpreter.execute( - "print('partial_timeout_state' in globals())\nprint('still alive')" - ) - finally: - interpreter.shutdown() - - assert "[Timeout] Iteration execution timed out after 0.2s" in timeout_result - assert "[stdout]\nbefore timeout" in timeout_result - assert "[stderr]\nstderr before timeout" in timeout_result - assert followup.strip() == "True\nstill alive" - - def test_predict_rlm_lm_selected_timeout_recovers_and_continues(self): - from predict_rlm import PredictRLM - - actions = SequentialActions( - SimpleNamespace( - reasoning="select a short timeout for a risky loop", - code=( - "import sys\n" - "print('before rlm timeout')\n" - "print('stderr before rlm timeout', file=sys.stderr)\n" - "while True:\n" - " pass\n" - ), - execution_timeout_seconds=0.2, - ), - SimpleNamespace( - reasoning="continue after the timeout observation", - code="SUBMIT(answer='continued after timeout')", - ), - ) - pool = SbxPool( - size=1, - config=SbxConfig(name=f"predict-rlm-test-rlm-timeout-{os.getpid()}"), - preinstall_packages=False, - ) - rlm = PredictRLM( - "prompt -> answer", - max_iterations=2, - sandbox_backend="sbx", - sbx_pool=pool, - ) - rlm.generate_action = actions - try: - prediction = rlm(prompt="exercise per-iteration timeout") - finally: - pool.shutdown() - - assert prediction.answer == "continued after timeout" - assert [call["iteration"] for call in actions.calls] == ["1/2", "2/2"] - assert len(prediction.trace.steps) == 2 - timeout_step, final_step = prediction.trace.steps - assert ( - "[Timeout] Iteration execution timed out after 0.2s" - in timeout_step.untruncated_output - ) - assert "[stdout]\nbefore rlm timeout" in timeout_step.untruncated_output - assert "[stderr]\nstderr before rlm timeout" in timeout_step.untruncated_output - assert final_step.output == "FINAL: {'answer': 'continued after timeout'}" - - def test_predict_rlm_can_use_predict_after_lm_selected_timeout(self): - from predict_rlm import PredictRLM - - actions = SequentialActions( - SimpleNamespace( - reasoning="call predict before the risky loop", - code=( - "first = await predict('question: str -> answer: str', " - "question='first call')\n" - "print('first predict:', first['answer'])\n" - "while True:\n" - " pass\n" - ), - execution_timeout_seconds=0.2, - ), - SimpleNamespace( - reasoning="call predict again after timeout recovery", - code=( - "second = await predict('question: str -> answer: str', " - "question='second call')\n" - "SUBMIT(answer=second['answer'])" - ), - ), - ) - class PredictionStub: - def __init__(self, answer: str) -> None: - self.answer = answer - - def keys(self) -> list[str]: - return ["answer"] - - def __getitem__(self, key: str) -> str: - return getattr(self, key) - - mock_lm = MagicMock() - mock_predictor = MagicMock() - mock_predictor.acall = AsyncMock(side_effect=[ - PredictionStub("pre-timeout prediction"), - PredictionStub("post-timeout prediction"), - ]) - pool = SbxPool( - size=1, - config=SbxConfig(name=f"predict-rlm-test-predict-timeout-{os.getpid()}", exec_timeout=12.0), - preinstall_packages=False, - ) - rlm = PredictRLM( - "prompt -> answer", - sub_lm=mock_lm, - max_iterations=2, - sandbox_backend="sbx", - sbx_pool=pool, - ) - rlm.generate_action = actions - try: - with patch("predict_rlm.predict_rlm.dspy.Predict", return_value=mock_predictor): - prediction = rlm(prompt="exercise predict across timeout recovery") - finally: - pool.shutdown() - - assert prediction.answer == "post-timeout prediction" - assert mock_predictor.acall.await_count == 2 - assert [call.kwargs["question"] for call in mock_predictor.acall.await_args_list] == [ - "first call", - "second call", - ] - timeout_step, final_step = prediction.trace.steps - assert "[Timeout] Iteration execution timed out after 0.2s" in timeout_step.untruncated_output - assert "first predict: pre-timeout prediction" in timeout_step.untruncated_output - assert final_step.output == "FINAL: {'answer': 'post-timeout prediction'}" - - def test_predict_rlm_recovers_after_user_exceptions_and_tools_still_work(self): - pool = SbxPool( - size=1, - config=SbxConfig(name=f"predict-rlm-test-user-exceptions-{os.getpid()}", exec_timeout=12.0), - preinstall_packages=False, - ) - try: - assert_predict_rlm_recovers_after_user_exceptions_and_tools_still_work(pool) - finally: - pool.shutdown() - - -class TestSbxBackendReattachConfig: - def test_reuse_requires_name(self): - with pytest.raises(ValidationError, match="reuse=True"): - SbxConfig(reuse=True) - - def test_reuse_implies_persist_and_no_remove(self): - config = SbxConfig(name="hot-box", reuse=True) - assert config.reuse is True - assert config.persist is True - assert config.remove_on_shutdown is False - - def test_reuse_false_is_unchanged_default(self): - config = SbxConfig() - assert config.reuse is False - assert config.persist is False - assert config.remove_on_shutdown is True - assert config.stop_on_shutdown is False - - -class TestSbxBackendReattachStagingRoot: - def test_reuse_staging_root_is_deterministic_from_name(self, tmp_path: Path): - with patch( - "predict_rlm.backends.sbx.backend.Path.cwd", return_value=tmp_path - ): - backend_a = SbxBackend(config=SbxConfig(name="hot-box", reuse=True)) - backend_b = SbxBackend(config=SbxConfig(name="hot-box", reuse=True)) - assert backend_a._staging_root == backend_b._staging_root - assert backend_a._staging_root.name == "hot-box" - - def test_reuse_staging_root_not_marked_for_cleanup(self, tmp_path: Path): - from predict_rlm.backends.sbx import backend as backend_mod - - with patch( - "predict_rlm.backends.sbx.backend.Path.cwd", return_value=tmp_path - ): - backend = SbxBackend(config=SbxConfig(name="hot-box", reuse=True)) - assert ( - str(backend._staging_root) - not in backend_mod._owned_staging_roots_pending_cleanup - ) - - def test_ephemeral_staging_root_is_unique_uuid(self, tmp_path: Path): - with patch( - "predict_rlm.backends.sbx.backend.Path.cwd", return_value=tmp_path - ): - backend_a = SbxBackend(config=SbxConfig()) - backend_b = SbxBackend(config=SbxConfig()) - assert backend_a._staging_root != backend_b._staging_root - - def test_reuse_relocated_staging_root_is_deterministic_across_sessions( - self, tmp_path: Path - ): - mounts = [DirectWorkspaceMount(host_path=str(tmp_path), sandbox_path="/work")] - - def _make() -> SbxBackend: - with patch( - "predict_rlm.backends.sbx.backend.Path.cwd", return_value=tmp_path - ): - return SbxBackend( - config=SbxConfig(name="hot-box", reuse=True), - direct_workspace_mounts=mounts, - ) - - backend_a = _make() - backend_b = _make() - try: - assert tmp_path not in backend_a._staging_root.parents - assert backend_a._staging_root == backend_b._staging_root - assert backend_a._staging_root.name == "predict-rlm-sbx-hot-box" - finally: - for backend in (backend_a, backend_b): - shutil.rmtree(backend._staging_root, ignore_errors=True) - - def test_ephemeral_relocated_staging_root_stays_unique(self, tmp_path: Path): - mounts = [DirectWorkspaceMount(host_path=str(tmp_path), sandbox_path="/work")] - with patch( - "predict_rlm.backends.sbx.backend.Path.cwd", return_value=tmp_path - ): - backend_a = SbxBackend(config=SbxConfig(), direct_workspace_mounts=mounts) - backend_b = SbxBackend(config=SbxConfig(), direct_workspace_mounts=mounts) - try: - assert tmp_path not in backend_a._staging_root.parents - assert backend_a._staging_root != backend_b._staging_root - finally: - for backend in (backend_a, backend_b): - shutil.rmtree(backend._staging_root, ignore_errors=True) - - -def _reattach_backend(tmp_path: Path, *, name: str = "hot-box") -> SbxBackend: - return SbxBackend( - config=SbxConfig(name=name, reuse=True), - preinstall_packages=False, - _staging_root=tmp_path / "staging", - ) - - -class TestSbxBackendReattachDetection: - def _patches(self, backend: SbxBackend, *, ls_output: str): - runs: list[list[str]] = [] - - def fake_run(cmd, *args, **kwargs): - runs.append(list(cmd)) - if cmd[:2] == ["sbx", "ls"]: - return SimpleNamespace(returncode=0, stdout=ls_output, stderr="") - return SimpleNamespace(returncode=0, stdout="", stderr="") - - cm = [ - patch( - "predict_rlm.backends.sbx.backend.shutil.which", - return_value="/usr/bin/sbx", - ), - patch( - "predict_rlm.backends.sbx.backend.subprocess.run", - side_effect=fake_run, - ), - patch.object( - SbxBackend, "_prepare_supervisor_script", return_value=Path("/sup.py") - ), - ] - return runs, cm - - def test_running_named_sandbox_reattaches_without_create_or_bootstrap( - self, tmp_path: Path - ): - backend = _reattach_backend(tmp_path) - runs, cms = self._patches(backend, ls_output="hot-box running\n") - with ( - cms[0], - cms[1], - cms[2], - patch.object(SbxBackend, "_apply_network_policy") as net, - patch.object(SbxBackend, "_bootstrap_packages") as boot, - patch.object(SbxBackend, "_setup_direct_workspace_aliases_in_sandbox"), - patch.object(SbxBackend, "_sbx_sandbox_healthy", return_value=True), - ): - backend._start_sbx_and_prepare_supervisor() - assert backend._sandbox_name == "hot-box" - assert not any(r[:2] == ["sbx", "create"] for r in runs) - net.assert_not_called() - boot.assert_not_called() - - def test_stopped_named_sandbox_is_started_then_reattaches(self, tmp_path: Path): - backend = _reattach_backend(tmp_path) - runs, cms = self._patches(backend, ls_output="hot-box stopped\n") - with ( - cms[0], - cms[1], - cms[2], - patch.object(SbxBackend, "_apply_network_policy") as net, - patch.object(SbxBackend, "_bootstrap_packages") as boot, - patch.object(SbxBackend, "_setup_direct_workspace_aliases_in_sandbox"), - patch.object(SbxBackend, "_sbx_sandbox_healthy", return_value=True), - ): - backend._start_sbx_and_prepare_supervisor() - assert backend._sandbox_name == "hot-box" - assert any( - r[:2] == ["sbx", "start"] and "hot-box" in r for r in runs - ), runs - assert not any(r[:2] == ["sbx", "create"] for r in runs) - net.assert_not_called() - boot.assert_not_called() - - def test_missing_named_sandbox_falls_through_to_create(self, tmp_path: Path): - backend = _reattach_backend(tmp_path) - runs, cms = self._patches(backend, ls_output="other-box running\n") - with ( - cms[0], - cms[1], - cms[2], - patch.object(SbxBackend, "_apply_network_policy") as net, - patch.object(SbxBackend, "_bootstrap_packages") as boot, - patch.object(SbxBackend, "_setup_direct_workspace_aliases_in_sandbox"), - patch.object(SbxBackend, "_sbx_sandbox_healthy", return_value=True), - ): - backend._start_sbx_and_prepare_supervisor() - assert backend._sandbox_name == "hot-box" - assert any(r[:2] == ["sbx", "create"] for r in runs), runs - net.assert_called_once() - boot.assert_called_once() - - def test_running_but_unhealthy_recreates(self, tmp_path: Path): - backend = _reattach_backend(tmp_path) - runs, _ = self._patches(backend, ls_output="hot-box running\n") - - def fake_run(cmd, *args, **kwargs): - runs.append(list(cmd)) - if cmd[:2] == ["sbx", "ls"]: - return SimpleNamespace( - returncode=0, stdout="hot-box running\n", stderr="" - ) - return SimpleNamespace(returncode=0, stdout="", stderr="") - - runs.clear() - with ( - patch( - "predict_rlm.backends.sbx.backend.shutil.which", - return_value="/usr/bin/sbx", - ), - patch( - "predict_rlm.backends.sbx.backend.subprocess.run", - side_effect=fake_run, - ), - patch.object( - SbxBackend, "_prepare_supervisor_script", return_value=Path("/sup.py") - ), - patch.object(SbxBackend, "_apply_network_policy") as net, - patch.object(SbxBackend, "_bootstrap_packages") as boot, - patch.object(SbxBackend, "_setup_direct_workspace_aliases_in_sandbox"), - patch.object(SbxBackend, "_sbx_sandbox_healthy", return_value=False), - ): - backend._start_sbx_and_prepare_supervisor() - assert any( - r[:2] == ["sbx", "rm"] and "hot-box" in r for r in runs - ), runs - assert any(r[:2] == ["sbx", "create"] for r in runs), runs - net.assert_called_once() - boot.assert_called_once() - - -class TestSbxBackendReattachShutdown: - def test_reuse_shutdown_does_not_rm_or_delete_staging(self, tmp_path: Path): - backend = _reattach_backend(tmp_path) - backend._sandbox_name = "hot-box" - staging = backend._staging_root - assert staging.exists() - runs: list[list[str]] = [] - - def fake_run(cmd, *args, **kwargs): - runs.append(list(cmd)) - return SimpleNamespace(returncode=0, stdout="", stderr="") - - with patch( - "predict_rlm.backends.sbx.backend.subprocess.run", side_effect=fake_run - ): - backend.shutdown() - assert not any(r[:2] == ["sbx", "rm"] for r in runs), runs - assert staging.exists() - - def test_reuse_stop_on_shutdown_stops_container(self, tmp_path: Path): - backend = SbxBackend( - config=SbxConfig(name="hot-box", reuse=True, stop_on_shutdown=True), - preinstall_packages=False, - _staging_root=tmp_path / "staging", - ) - backend._sandbox_name = "hot-box" - runs: list[list[str]] = [] - - def fake_run(cmd, *args, **kwargs): - runs.append(list(cmd)) - return SimpleNamespace(returncode=0, stdout="", stderr="") - - with patch( - "predict_rlm.backends.sbx.backend.subprocess.run", side_effect=fake_run - ): - backend.shutdown() - assert any( - r[:2] == ["sbx", "stop"] and "hot-box" in r for r in runs - ), runs - assert not any(r[:2] == ["sbx", "rm"] for r in runs), runs - - -class TestSbxBackendDestroy: - def test_destroy_removes_sandbox_and_staging_root(self, tmp_path: Path): - backend = _reattach_backend(tmp_path) - backend._sandbox_name = "hot-box" - staging = backend._staging_root - assert staging.exists() - runs: list[list[str]] = [] - - def fake_run(cmd, *args, **kwargs): - runs.append(list(cmd)) - return SimpleNamespace(returncode=0, stdout="", stderr="") - - with patch( - "predict_rlm.backends.sbx.backend.subprocess.run", side_effect=fake_run - ): - backend.destroy() - assert any( - r[:3] == ["sbx", "rm", "--force"] and "hot-box" in r for r in runs - ), runs - assert not staging.exists() - - def test_remove_classmethod_force_removes_named_sandbox(self): - runs: list[list[str]] = [] - - def fake_run(cmd, *args, **kwargs): - runs.append(list(cmd)) - return SimpleNamespace(returncode=0, stdout="", stderr="") - - with patch( - "predict_rlm.backends.sbx.backend.subprocess.run", side_effect=fake_run - ): - SbxBackend.remove("hot-box") - assert any( - r[:3] == ["sbx", "rm", "--force"] and "hot-box" in r for r in runs - ), runs - - -class TestSbxBackendReattachRegression: - def test_default_path_still_creates_without_ls_probe(self, tmp_path: Path): - backend = SbxBackend( - config=SbxConfig(), - preinstall_packages=False, - _staging_root=tmp_path / "staging", - ) - runs: list[list[str]] = [] - - def fake_run(cmd, *args, **kwargs): - runs.append(list(cmd)) - return SimpleNamespace(returncode=0, stdout="auto-name\n", stderr="") - - with ( - patch( - "predict_rlm.backends.sbx.backend.shutil.which", - return_value="/usr/bin/sbx", - ), - patch( - "predict_rlm.backends.sbx.backend.subprocess.run", side_effect=fake_run - ), - patch.object( - SbxBackend, "_prepare_supervisor_script", return_value=Path("/sup.py") - ), - patch.object(SbxBackend, "_apply_network_policy") as net, - patch.object(SbxBackend, "_bootstrap_packages") as boot, - patch.object(SbxBackend, "_setup_direct_workspace_aliases_in_sandbox"), - ): - backend._start_sbx_and_prepare_supervisor() - assert not any(r[:2] == ["sbx", "ls"] for r in runs), runs - assert any(r[:2] == ["sbx", "create"] for r in runs), runs - net.assert_called_once() - boot.assert_called_once() - - -@pytest.mark.integration -@pytest.mark.skipif( - not _real_sbx_available(), - reason="real Docker Sandboxes tests require PREDICT_RLM_RUN_SBX_TESTS=1, sbx CLI, and sbx login", -) -class TestSbxBackendRealSbxReattach: - def _list_names(self) -> list[str]: - result = subprocess.run( - ["sbx", "ls"], capture_output=True, text=True, check=False, timeout=15 - ) - return [line.split()[0] for line in result.stdout.splitlines() if line.split()] - - def test_persist_reattach_destroy_lifecycle(self): - name = f"predict-rlm-reattach-{os.getpid()}" - config = SbxConfig(name=name, reuse=True) - marker = f"state-{os.getpid()}" - - first = SbxBackend(config=config, preinstall_packages=False, debug=True) - try: - first.prewarm() - first.execute( - "from pathlib import Path\n" - f"Path('/sandbox/persisted.txt').write_text({marker!r})\n" - "print('wrote')" - ) - first.shutdown() - assert name in self._list_names() - - second = SbxBackend(config=config, preinstall_packages=False, debug=True) - events: list[str] = [] - orig_log = second._log_lifecycle - - def spy_log(event, **fields): - events.append(event) - return orig_log(event, **fields) - - with ( - patch.object(second, "_log_lifecycle", side_effect=spy_log), - patch.object( - SbxBackend, - "_bootstrap_packages", - side_effect=AssertionError("bootstrap must not run on reattach"), - ), - ): - second.prewarm() - out = second.execute( - "from pathlib import Path\n" - "print(Path('/sandbox/persisted.txt').read_text())" - ) - assert out.strip() == marker - assert any(e.startswith("sbx.reattach") for e in events), events - assert not any(e == "sbx.create.start" for e in events), events - second.shutdown() - assert name in self._list_names() - - second.destroy() - assert name not in self._list_names() + second.destroy() + assert name not in self._list_names() third = SbxBackend(config=config, preinstall_packages=False, debug=True) try: third.prewarm() fresh = third.execute( - "from pathlib import Path\n" - "print(Path('/sandbox/persisted.txt').exists())" + "from pathlib import Path\nprint(Path('/sandbox/persisted.txt').exists())" ) assert fresh.strip() == "False" finally: @@ -4461,33 +1418,3 @@ def spy_log(event, **fields): text=True, check=False, ) - - def test_reattach_after_interpreter_error_recovers(self): - name = f"predict-rlm-recover-{os.getpid()}" - config = SbxConfig(name=name, reuse=True) - - first = SbxBackend(config=config, preinstall_packages=False, debug=True) - try: - first.prewarm() - first.execute("keep = 7\nprint('ready')") - with pytest.raises(CodeInterpreterError, match="ValueError"): - first.execute("raise ValueError('boom')") - assert first.execute("print(keep + 1)").strip() == "8" - first.shutdown() - assert name in self._list_names() - - second = SbxBackend(config=config, preinstall_packages=False, debug=True) - second.prewarm() - assert second.execute("print('recovered')").strip() == "recovered" - with pytest.raises(CodeInterpreterError, match="ValueError"): - second.execute("raise ValueError('again')") - assert second.execute("print(6 * 7)").strip() == "42" - second.destroy() - assert name not in self._list_names() - finally: - subprocess.run( - ["sbx", "rm", "--force", name], - capture_output=True, - text=True, - check=False, - ) diff --git a/tests/test_sbx_pool.py b/tests/test_sbx_pool.py index 755f5781..c1efa432 100644 --- a/tests/test_sbx_pool.py +++ b/tests/test_sbx_pool.py @@ -3,6 +3,9 @@ from __future__ import annotations import asyncio +import sys +import threading +import time from contextlib import asynccontextmanager from pathlib import Path from types import SimpleNamespace @@ -12,7 +15,7 @@ pytest.importorskip("websockets") -from predict_rlm.backends.sbx import SbxPool # noqa: E402 +from predict_rlm.backends.sbx import SbxConfig, SbxPool # noqa: E402 from predict_rlm.backends.sbx.execution import SbxPoolExecutionBackend # noqa: E402 from predict_rlm.runtime import ( # noqa: E402 ExecutionSpec, @@ -145,41 +148,6 @@ def create_interpreter(index: int) -> SyncFakeInterpreter: return pool, events, created -@pytest.mark.asyncio -async def test_alease_is_exclusive_and_awaits_async_lifecycle(tmp_path: Path, monkeypatch): - pool, events, created = make_pool(tmp_path, monkeypatch) - second_acquired = asyncio.Event() - - def forbidden_to_thread(*args: Any, **kwargs: Any) -> None: - raise AssertionError("async SBX pool lifecycle delegated to a thread") - - monkeypatch.setattr(asyncio, "to_thread", forbidden_to_thread) - - async def second_lease() -> None: - async with pool.alease() as interpreter: - assert interpreter is created[0] - second_acquired.set() - - async with pool.alease(tools={"tool": lambda: None}) as interpreter: - assert interpreter is created[0] - waiter = asyncio.create_task(second_lease()) - await asyncio.sleep(0) - assert not second_acquired.is_set() - - await waiter - await pool.ashutdown() - - assert [event[0] for event in events] == [ - "prewarm", - "configure", - "reset", - "configure", - "reset", - "shutdown", - ] - assert list(events[1][2]["tools"]) == ["tool"] - - @pytest.mark.asyncio async def test_alease_replaces_interpreter_after_reset_failure(tmp_path: Path, monkeypatch): pool, events, created = make_pool(tmp_path, monkeypatch) @@ -197,61 +165,6 @@ async def test_alease_replaces_interpreter_after_reset_failure(tmp_path: Path, m await pool.ashutdown() -@pytest.mark.asyncio -async def test_alease_retires_busy_interpreter_without_reset_or_immediate_shutdown( - tmp_path: Path, - monkeypatch, -): - pool, events, created = make_pool(tmp_path, monkeypatch) - - async with pool.alease() as interpreter: - interpreter.live_host_work = True - interpreter.fail_reset = True - - assert len(created) == 2 - assert ("reset", 0) not in events - assert ("aretire", 0) in events - assert ("shutdown", 0) in events - assert list(pool._available.queue) == [created[1]] - - created[0].live_host_work = False - await pool.ashutdown() - - -@pytest.mark.asyncio -async def test_alease_awaits_busy_interpreter_retirement_before_replacement( - tmp_path: Path, - monkeypatch, -): - pool, events, created = make_pool(tmp_path, monkeypatch) - lease_entered = asyncio.Event() - retirement_started = asyncio.Event() - retirement_release = asyncio.Event() - - async def use_busy_interpreter() -> None: - async with pool.alease() as interpreter: - interpreter.live_host_work = True - interpreter.retirement_started = retirement_started - interpreter.retirement_release = retirement_release - lease_entered.set() - - lease = asyncio.create_task(use_busy_interpreter()) - await lease_entered.wait() - await asyncio.sleep(0.02) - - try: - assert retirement_started.is_set() - assert not lease.done() - assert pool._available.empty() - finally: - retirement_release.set() - await lease - - assert ("shutdown", 0) in events - assert list(pool._available.queue) == [created[1]] - await pool.ashutdown() - - @pytest.mark.asyncio async def test_cancelled_alease_still_finishes_busy_interpreter_retirement( tmp_path: Path, @@ -419,7 +332,9 @@ def create_failed_replacement(index): @pytest.mark.asyncio -async def test_alease_releases_interpreter_when_configuration_fails(tmp_path: Path, monkeypatch): +async def test_alease_releases_interpreter_when_configuration_fails( + tmp_path: Path, monkeypatch +): pool, events, created = make_pool(tmp_path, monkeypatch) await pool.astart() created[0].fail_configure = True @@ -601,18 +516,249 @@ async def alease(self, **kwargs): assert pool.acquisitions == 1 -def test_pool_exposes_immutable_fixed_session_requirements(tmp_path: Path): - pool = SbxPool( - size=1, - allowed_domains=["service.internal"], - extra_read_paths=["/host/input"], - extra_write_paths=["/host/output"], - preinstall_packages=False, - _staging_root=tmp_path / "policy-pool", - ) +PAYLOAD_PATH = Path(__file__).parents[1] / "src/predict_rlm/backends/supervisor/_payload.py" - assert pool.session_requirements == SessionRequirements( - allowed_domains=("service.internal",), - extra_read_paths=("/host/input",), - extra_write_paths=("/host/output",), - ) + +@pytest.mark.sbx +class TestSbxPool: + def test_start_failure_shuts_down_created_interpreters_and_leaves_pool_stopped( + self, tmp_path: Path, monkeypatch + ): + pool = SbxPool( + size=3, + config=SbxConfig(name="pool-test"), + preinstall_packages=False, + _staging_root=tmp_path / "pool", + ) + created = [] + + class FakeInterpreter: + def __init__(self, index: int) -> None: + self.index = index + self.shutdown_called = False + + def prewarm(self) -> None: + if self.index == 1: + raise RuntimeError("prewarm failed") + + def shutdown(self) -> None: + self.shutdown_called = True + + def create_interpreter(index: int) -> FakeInterpreter: + interpreter = FakeInterpreter(index) + created.append(interpreter) + return interpreter + + monkeypatch.setattr(pool, "_create_interpreter", create_interpreter) + + with pytest.raises(RuntimeError, match="prewarm failed"): + pool.start() + + assert created + assert all(interpreter.shutdown_called for interpreter in created) + assert not pool._started + assert pool._all_interpreters == [] + assert pool._available.qsize() == 0 + + def test_shutdown_runs_concurrently_and_attempts_all_interpreters( + self, tmp_path: Path, monkeypatch + ): + pool = SbxPool( + size=3, + config=SbxConfig(name="pool-test"), + preinstall_packages=False, + _staging_root=tmp_path / "pool", + ) + barrier = threading.Barrier(3) + active = 0 + max_active = 0 + active_lock = threading.Lock() + shutdown_indexes: list[int] = [] + + class FakeInterpreter: + def __init__(self, index: int) -> None: + self.index = index + + def prewarm(self) -> None: + return None + + def shutdown(self) -> None: + nonlocal active, max_active + with active_lock: + active += 1 + max_active = max(max_active, active) + try: + barrier.wait(timeout=1) + shutdown_indexes.append(self.index) + if self.index == 1: + raise RuntimeError("shutdown failed") + finally: + with active_lock: + active -= 1 + + monkeypatch.setattr(pool, "_create_interpreter", lambda index: FakeInterpreter(index)) + pool.start() + + with pytest.raises(RuntimeError, match="shutdown failed"): + pool.shutdown() + + assert max_active == 3 + assert sorted(shutdown_indexes) == [0, 1, 2] + assert not pool._started + assert pool._shutdown + assert pool._all_interpreters == [] + assert pool._available.qsize() == 0 + + def test_lease_is_exclusive_and_release_resets(self, tmp_path: Path): + pool = SbxPool( + size=1, + config=SbxConfig(name="pool-test"), + preinstall_packages=False, + _supervisor_command=[sys.executable, "-u", str(PAYLOAD_PATH)], + _staging_root=tmp_path / "pool", + ) + acquired = threading.Event() + released = threading.Event() + + def second_lease() -> None: + with pool.lease() as interpreter: + acquired.set() + assert interpreter.execute("print('x' in globals())").strip() == "False" + + try: + pool.start() + with pool.lease() as interpreter: + interpreter.execute("x = 7") + staged = pool._all_interpreters[0]._host_path_for_virtual_path( + "/sandbox/output/value.txt" + ) + staged.parent.mkdir(parents=True, exist_ok=True) + staged.write_text("leaked", encoding="utf-8") + thread = threading.Thread(target=second_lease) + thread.start() + assert not acquired.wait(0.2) + + released.set() + thread.join(timeout=5) + + assert released.is_set() + assert acquired.is_set() + with pool.lease() as interpreter: + assert interpreter.list_dir("/sandbox") == [] + finally: + pool.shutdown() + + def test_shutdown_requested_during_start_prevents_waiting_lease_acquire( + self, tmp_path: Path, monkeypatch + ): + pool = SbxPool( + size=1, + config=SbxConfig(name="pool-test"), + preinstall_packages=False, + _staging_root=tmp_path / "pool", + ) + prewarm_started = threading.Event() + allow_prewarm = threading.Event() + + class FakeInterpreter: + def __init__(self) -> None: + self.shutdown_called = False + + def prewarm(self) -> None: + prewarm_started.set() + assert allow_prewarm.wait(timeout=2) + + def configure_runtime(self, **kwargs) -> None: + return None + + def reset(self) -> None: + return None + + def shutdown(self) -> None: + self.shutdown_called = True + + interpreter = FakeInterpreter() + monkeypatch.setattr(pool, "_create_interpreter", lambda index: interpreter) + + lease_results: list[str] = [] + + def lease_during_start() -> None: + try: + with pool.lease(): + lease_results.append("acquired") + except RuntimeError as exc: + lease_results.append(str(exc)) + + lease_thread = threading.Thread(target=lease_during_start) + lease_thread.start() + assert prewarm_started.wait(timeout=2) + + shutdown_thread = threading.Thread(target=pool.shutdown) + shutdown_thread.start() + deadline = time.monotonic() + 2 + while time.monotonic() < deadline: + with pool._state_changed: + if pool._shutdown_requested: + break + time.sleep(0.01) + else: + pytest.fail("shutdown did not request pool stop while startup was active") + allow_prewarm.set() + + lease_thread.join(timeout=2) + shutdown_thread.join(timeout=2) + + assert not lease_thread.is_alive() + assert not shutdown_thread.is_alive() + assert lease_results == ["SbxPool is shut down"] + assert interpreter.shutdown_called + assert pool._available.qsize() == 0 + + def test_lease_after_shutdown_raises_until_explicit_restart( + self, tmp_path: Path, monkeypatch + ): + pool = SbxPool( + size=1, + config=SbxConfig(name="pool-test"), + preinstall_packages=False, + _staging_root=tmp_path / "pool", + ) + + class FakeInterpreter: + def __init__(self, index: int) -> None: + self.index = index + + def prewarm(self) -> None: + return None + + def configure_runtime(self, **kwargs) -> None: + return None + + def reset(self) -> None: + return None + + def shutdown(self) -> None: + return None + + created: list[FakeInterpreter] = [] + + def create_interpreter(index: int) -> FakeInterpreter: + interpreter = FakeInterpreter(index) + created.append(interpreter) + return interpreter + + monkeypatch.setattr(pool, "_create_interpreter", create_interpreter) + + pool.start() + pool.shutdown() + + with pytest.raises(RuntimeError, match="SbxPool is shut down"): + with pool.lease(): + pass + + pool.start() + try: + with pool.lease() as interpreter: + assert interpreter is created[-1] + finally: + pool.shutdown() diff --git a/tests/test_shared.py b/tests/test_shared.py deleted file mode 100644 index 4f95a8d0..00000000 --- a/tests/test_shared.py +++ /dev/null @@ -1,209 +0,0 @@ -"""Tests for _shared.py: format_tool_docs_full and build_rlm_signatures.""" - -from unittest.mock import MagicMock, patch - -import dspy - -from predict_rlm._shared import ( - build_rlm_signatures, - format_tool_docs_full, - strip_code_fences, -) - - -def test_strip_code_fences_accepts_supported_fences_and_bare_code(): - assert strip_code_fences("```python\nprint('python')\n```") == "print('python')" - assert strip_code_fences("```py\nprint('py')\n```") == "print('py')" - assert strip_code_fences("```repl\nprint('repl')\n```") == "print('repl')" - assert strip_code_fences("```\nprint('bare fence')\n```") == "print('bare fence')" - assert strip_code_fences("print('bare code')") == "print('bare code')" - - -class TestFormatToolDocsFull: - def test_empty_tools_returns_empty_string(self): - assert format_tool_docs_full({}) == "" - - def test_single_tool_with_signature_and_docstring(self): - def fetch_page(url: str, timeout: int) -> str: - """Fetch a web page. - - Args: - url: The URL to fetch. - timeout: Request timeout in seconds. - - Returns: - The page HTML content. - """ - return "" - - result = format_tool_docs_full({"fetch_page": fetch_page}) - assert "### `await fetch_page(url: str, timeout: int) -> str`" in result - assert "Fetch a web page." in result - assert "Args:" in result - assert "Returns:" in result - - def test_tool_without_docstring(self): - def no_docs(x: str) -> str: - return x - - no_docs.__doc__ = None - result = format_tool_docs_full({"no_docs": no_docs}) - assert "No description" in result - - def test_tool_with_uninspectable_signature(self): - mock_fn = MagicMock(spec=lambda: None) - # Override signature inspection to raise - mock_fn.__doc__ = "A mock tool." - with patch("predict_rlm._shared.inspect.signature", side_effect=ValueError): - result = format_tool_docs_full({"mock_fn": mock_fn}) - assert "mock_fn(...)" in result - assert "A mock tool." in result - - def test_multiple_tools_all_listed(self): - def tool_a() -> str: - """Tool A.""" - return "" - - def tool_b() -> str: - """Tool B.""" - return "" - - result = format_tool_docs_full({"tool_a": tool_a, "tool_b": tool_b}) - assert "tool_a" in result - assert "tool_b" in result - assert "## Additional Tools" in result - - def test_header_mentions_async_and_gather(self): - def dummy() -> str: - """Dummy.""" - return "" - - result = format_tool_docs_full({"dummy": dummy}) - assert "await" in result - assert "asyncio.gather()" in result - - -class TestBuildRlmSignatures: - ACTION_TEMPLATE = ( - "You have inputs: {inputs}.\n" - "Output fields:\n{output_fields}\n" - "Submit: {final_output_names}" - ) - - def test_returns_action_and_extract_signatures(self): - sig = dspy.Signature("question -> answer") - action, extract = build_rlm_signatures( - sig, self.ACTION_TEMPLATE, {}, format_tool_docs_full - ) - assert "variables_info" in action.input_fields - assert "repl_history" in action.input_fields - assert "iteration" in action.input_fields - assert "reasoning" in action.output_fields - assert "code" in action.output_fields - - assert "variables_info" in extract.input_fields - assert "repl_history" in extract.input_fields - assert "answer" in extract.output_fields - - def test_action_signature_exposes_optional_execution_timeout(self): - sig = dspy.Signature("question -> answer") - action, _ = build_rlm_signatures( - sig, self.ACTION_TEMPLATE, {}, format_tool_docs_full - ) - - timeout_field = action.output_fields.get("execution_timeout_seconds") - assert timeout_field is not None - assert timeout_field.default is None - assert not timeout_field.is_required() - assert type(None) in getattr(timeout_field.annotation, "__args__", ()) - desc = timeout_field.json_schema_extra["desc"] - assert "Use null for ordinary short, safe code" in desc - assert "positive number of seconds" in desc - assert "loops" in desc - assert "large scans" in desc - assert "tool/network fanout" in desc - assert "batch predict() calls" in desc - assert "tests/subprocesses" in desc - assert "stdout/stderr" in desc - assert "next iteration can continue" in desc - - def test_action_signature_omits_execution_timeout_when_disabled(self): - sig = dspy.Signature("question -> answer") - action, _ = build_rlm_signatures( - sig, - self.ACTION_TEMPLATE, - {}, - format_tool_docs_full, - model_execution_timeout=False, - ) - - assert "execution_timeout_seconds" not in action.output_fields - assert "code" in action.output_fields # other output fields unaffected - - def test_action_signature_keeps_code_field_description_narrow(self): - sig = dspy.Signature("question -> answer") - action, _ = build_rlm_signatures( - sig, self.ACTION_TEMPLATE, {}, format_tool_docs_full - ) - - assert ( - action.output_fields["code"].json_schema_extra["desc"] - == "Python code wrapped in ```repl blocks." - ) - - def test_tool_docs_in_action_instructions(self): - def my_tool(x: str) -> str: - """Does something useful.""" - return x - - sig = dspy.Signature("question -> answer") - action, _ = build_rlm_signatures( - sig, self.ACTION_TEMPLATE, {"my_tool": my_tool}, format_tool_docs_full - ) - assert "my_tool" in action.instructions - - def test_skill_instructions_appended(self): - sig = dspy.Signature("question -> answer") - action, _ = build_rlm_signatures( - sig, - self.ACTION_TEMPLATE, - {}, - format_tool_docs_full, - skill_instructions="Use pdfplumber for PDF extraction.", - ) - assert "## Skills" in action.instructions - assert "Use pdfplumber for PDF extraction." in action.instructions - - def test_file_instructions_appended(self): - sig = dspy.Signature("question -> answer") - action, _ = build_rlm_signatures( - sig, - self.ACTION_TEMPLATE, - {}, - format_tool_docs_full, - file_instructions="Input files are mounted at /sandbox/input.", - ) - assert "Input files are mounted at /sandbox/input." in action.instructions - - def test_original_signature_instructions_preserved(self): - sig = dspy.Signature("question -> answer", "Be concise and accurate.") - action, extract = build_rlm_signatures( - sig, self.ACTION_TEMPLATE, {}, format_tool_docs_full - ) - assert "Be concise and accurate." in action.instructions - assert "Be concise and accurate." in extract.instructions - - def test_no_skill_or_file_instructions(self): - sig = dspy.Signature("question -> answer") - action, _ = build_rlm_signatures( - sig, self.ACTION_TEMPLATE, {}, format_tool_docs_full - ) - assert "## Skills" not in action.instructions - - def test_extract_sig_includes_output_fields(self): - sig = dspy.Signature("question -> answer, confidence: float") - _, extract = build_rlm_signatures( - sig, self.ACTION_TEMPLATE, {}, format_tool_docs_full - ) - assert "answer" in extract.output_fields - assert "confidence" in extract.output_fields diff --git a/tests/test_skill_package_integration.py b/tests/test_skill_package_integration.py index 15627d25..433bf81e 100644 --- a/tests/test_skill_package_integration.py +++ b/tests/test_skill_package_integration.py @@ -17,6 +17,7 @@ packages=["python-slugify"], ) + SLUGIFY_CODE = ( "from slugify import slugify\n" "answer = slugify('Skill Packages Work')\n" @@ -74,51 +75,12 @@ def _run_skill_package_rlm(**kwargs) -> None: assert len(actions.calls) == 1 -@pytest.mark.integration -@pytest.mark.skipif(shutil.which("deno") is None, reason="JSPI skill-package test requires Deno") -def test_predict_rlm_installs_skill_packages_in_jspi_sandbox() -> None: - _run_skill_package_rlm() - - -@pytest.mark.sbx -@pytest.mark.integration -@pytest.mark.skipif( - not _real_sbx_available(), - reason="real SBX tests require PREDICT_RLM_RUN_SBX_TESTS=1, sbx CLI, and sbx login", -) -def test_predict_rlm_installs_skill_packages_in_sbx_sandbox() -> None: - from predict_rlm import SbxConfig - - _run_skill_package_rlm( - sandbox_backend="sbx", - sbx_config=SbxConfig(name=f"predict-rlm-skill-package-{os.getpid()}"), - ) - - -@pytest.mark.sbx @pytest.mark.integration @pytest.mark.skipif( - not _real_sbx_available(), - reason="real SBX tests require PREDICT_RLM_RUN_SBX_TESTS=1, sbx CLI, and sbx login", + shutil.which("deno") is None, reason="JSPI skill-package test requires Deno" ) -def test_predict_rlm_installs_skill_packages_in_injected_sbx_sandbox(tmp_path) -> None: - from predict_rlm import SbxConfig - from predict_rlm.backends import SbxBackend - from predict_rlm.workspace import DirectWorkspaceMount - - interpreter = SbxBackend( - config=SbxConfig(name=f"predict-rlm-skill-package-injected-{os.getpid()}"), - direct_workspace_mounts=[ - DirectWorkspaceMount( - host_path=str(tmp_path.resolve()), - sandbox_path=str(tmp_path.resolve()), - ) - ], - ) - try: - _run_skill_package_rlm(interpreter=interpreter) - finally: - interpreter.shutdown() +def test_predict_rlm_installs_skill_packages_in_jspi_sandbox() -> None: + _run_skill_package_rlm() @pytest.mark.sbx diff --git a/tests/test_small_kernel.py b/tests/test_small_kernel.py index ef77160e..2704f1b7 100644 --- a/tests/test_small_kernel.py +++ b/tests/test_small_kernel.py @@ -1,10 +1,8 @@ from __future__ import annotations import asyncio -import inspect import shutil import threading -from collections.abc import Sequence from contextlib import asynccontextmanager, contextmanager from pathlib import Path from typing import Annotated @@ -18,121 +16,26 @@ InterpreterBackendAdapter, InterpreterExecutionSession, ) -from predict_rlm.compatibility import ( - FileInputAdapter, - FileOutputAdapter, - SyncedFileToolOperation, -) +from predict_rlm.compatibility import SyncedFileToolOperation from predict_rlm.evidence import ( EvidenceIncompleteError, EvidenceRecorder, RunEventKind, ) -from predict_rlm.files import File as RuntimeFile from predict_rlm.files import SyncedFile from predict_rlm.runtime import ( - Artifact, ArtifactBinding, BoundInput, CallableTool, ExecutionSpec, HostDirectoryMount, InputAdapter, - OutputAdapter, PreparedInput, RunContext, - RuntimeContribution, RuntimeSpec, SessionOwnership, - SessionRequirements, - callable_has_sync_leaf, - current_run_context, - resolve_runtime_spec, use_run_context, ) -from predict_rlm.workspace import Workspace - - -class KernelFileSignature(dspy.Signature): - source: RuntimeFile = dspy.InputField() - result: RuntimeFile = dspy.OutputField() - - -class FileUnionInputSignature(dspy.Signature): - source: RuntimeFile | str = dspy.InputField() - answer: str = dspy.OutputField() - - -class NestedFileInputSignature(dspy.Signature): - source: list[list[RuntimeFile]] = dspy.InputField() - answer: str = dspy.OutputField() - - -class WorkspaceUnionInputSignature(dspy.Signature): - source: Workspace | str = dspy.InputField() - answer: str = dspy.OutputField() - - -class NestedWorkspaceInputSignature(dspy.Signature): - source: list[list[Workspace]] = dspy.InputField() - answer: str = dspy.OutputField() - - -class TupleFileInputSignature(dspy.Signature): - source: tuple[RuntimeFile, ...] = dspy.InputField() - answer: str = dspy.OutputField() - - -class SetFileInputSignature(dspy.Signature): - source: set[RuntimeFile] = dspy.InputField() - answer: str = dspy.OutputField() - - -class DictFileInputSignature(dspy.Signature): - source: dict[str, RuntimeFile] = dspy.InputField() - answer: str = dspy.OutputField() - - -class SequenceFileInputSignature(dspy.Signature): - source: Sequence[RuntimeFile] = dspy.InputField() - answer: str = dspy.OutputField() - - -class NestedGenericFileInputSignature(dspy.Signature): - source: list[dict[str, tuple[RuntimeFile, ...]]] = dspy.InputField() - answer: str = dspy.OutputField() - - -class TupleWorkspaceInputSignature(dspy.Signature): - source: tuple[Workspace, ...] = dspy.InputField() - answer: str = dspy.OutputField() - - -class SetWorkspaceInputSignature(dspy.Signature): - source: set[Workspace] = dspy.InputField() - answer: str = dspy.OutputField() - - -class DictWorkspaceInputSignature(dspy.Signature): - source: dict[str, Workspace] = dspy.InputField() - answer: str = dspy.OutputField() - - -class SequenceWorkspaceInputSignature(dspy.Signature): - source: Sequence[Workspace] = dspy.InputField() - answer: str = dspy.OutputField() - - -class NestedGenericWorkspaceInputSignature(dspy.Signature): - source: list[dict[str, tuple[Workspace, ...]]] = dspy.InputField() - answer: str = dspy.OutputField() - - -class StubBackend: - name = "stub" - - def start(self, spec, ctx): - raise NotImplementedError def make_spec(*, events=()) -> RuntimeSpec: @@ -141,173 +44,11 @@ def make_spec(*, events=()) -> RuntimeSpec: adapters=(), tools=(), packages=(), - execution=StubBackend(), + execution=FinalBackend(), events=events, ) -def test_predict_rlm_exposes_one_adapter_extension_parameter(): - from predict_rlm import PredictRLM - - parameters = inspect.signature(PredictRLM.__init__).parameters - - assert "adapters" in parameters - assert parameters["events"].annotation == "Sequence[EventSink]" - assert "inputs" not in parameters - assert "outputs" not in parameters - - -def test_runtime_spec_partitions_an_immutable_adapter_snapshot_by_role(): - class SharedInputAdapter(InputAdapter[str]): - name = "shared" - value_type = str - - async def prepare(self, field, value, ctx): - return PreparedInput(model_value=value) - - class SharedOutputAdapter(OutputAdapter[str]): - name = "shared" - value_type = str - - async def reserve(self, field, value, ctx, session): - raise NotImplementedError - - async def materialize(self, reservation, submitted_value, ctx, session): - raise NotImplementedError - - input_adapter = SharedInputAdapter() - output_adapter = SharedOutputAdapter() - - def output_module(): - return RuntimeContribution(adapters=[output_adapter]) - - contribution = RuntimeContribution( - adapters=[input_adapter], - execution=StubBackend(), - ) - spec = resolve_runtime_spec(direct=contribution, modules=(output_module,)) - - assert isinstance(contribution.adapters, tuple) - assert isinstance(spec.adapters, tuple) - assert spec.adapters == (input_adapter, output_adapter) - assert spec.input_adapters == (input_adapter,) - assert spec.output_adapters == (output_adapter,) - - -def test_runtime_spec_rejects_duplicate_adapter_names_within_one_role(): - class DuplicateInputAdapter(InputAdapter[str]): - name = "duplicate" - value_type = str - - async def prepare(self, field, value, ctx): - return PreparedInput(model_value=value) - - with pytest.raises(ValueError, match="Duplicate input adapter name"): - resolve_runtime_spec( - direct=RuntimeContribution( - adapters=[DuplicateInputAdapter(), DuplicateInputAdapter()], - execution=StubBackend(), - ) - ) - - -def test_configured_adapter_names_suppress_compatibility_defaults_per_role(): - from predict_rlm import PredictRLM - - class CustomFileInputAdapter(InputAdapter[RuntimeFile]): - name = "file" - value_type = RuntimeFile - - async def prepare(self, field, value, ctx): - return PreparedInput(model_value=value) - - class CustomFileOutputAdapter(OutputAdapter[RuntimeFile]): - name = "file" - value_type = RuntimeFile - - async def reserve(self, field, value, ctx, session): - raise NotImplementedError - - async def materialize(self, reservation, submitted_value, ctx, session): - raise NotImplementedError - - custom_input = CustomFileInputAdapter() - input_override = PredictRLM(KernelFileSignature, adapters=[custom_input]) - assert custom_input in input_override.runtime_spec.input_adapters - assert not any( - isinstance(adapter, FileInputAdapter) - for adapter in input_override.runtime_spec.input_adapters - ) - assert any( - isinstance(adapter, FileOutputAdapter) - for adapter in input_override.runtime_spec.output_adapters - ) - - custom_output = CustomFileOutputAdapter() - output_override = PredictRLM(KernelFileSignature, adapters=[custom_output]) - assert any( - isinstance(adapter, FileInputAdapter) - for adapter in output_override.runtime_spec.input_adapters - ) - assert custom_output in output_override.runtime_spec.output_adapters - assert not any( - isinstance(adapter, FileOutputAdapter) - for adapter in output_override.runtime_spec.output_adapters - ) - - -def test_runtime_spec_expands_modules_once_and_deduplicates_packages(): - calls = 0 - - def module(): - nonlocal calls - calls += 1 - return RuntimeContribution(instructions=("module",), packages=("b", "a")) - - spec = resolve_runtime_spec( - direct=RuntimeContribution( - instructions=("direct",), - packages=("a",), - execution=StubBackend(), - ), - modules=(module,), - ) - - assert calls == 1 - assert spec.instructions == ("direct", "module") - assert spec.packages == ("a", "b") - - -def test_runtime_spec_rejects_duplicate_tool_names(): - first = CallableTool(name="same", function=lambda: 1) - second = CallableTool(name="same", function=lambda: 2) - - with pytest.raises(ValueError, match="Duplicate tool name"): - resolve_runtime_spec( - direct=RuntimeContribution( - tools=(first, second), - execution=StubBackend(), - ) - ) - - -@pytest.mark.asyncio -async def test_run_context_is_invocation_local(): - first = RunContext(make_spec(), {"value": 1}) - second = RunContext(make_spec(), {"value": 2}) - - async with use_run_context(first): - assert current_run_context() is first - first.state["marker"] = "first" - async with use_run_context(second): - assert current_run_context() is second - assert "marker" not in second.state - assert current_run_context() is first - - assert current_run_context() is None - assert first.run_id != second.run_id - - class RecordingSink: strict = True @@ -394,20 +135,6 @@ async def test_strict_evidence_close_failure_cannot_publish_success(): assert not ctx.evidence_complete -@pytest.mark.asyncio -async def test_evidence_failure_does_not_replace_primary_failure(): - sink = RecordingSink(fail_flush=True) - ctx = RunContext(make_spec(events=(sink,)), {}) - recorder = EvidenceRecorder(ctx, (sink,)) - primary = RuntimeError("primary") - - await recorder.emit(RunEventKind.RUN_STARTED) - await recorder.finish_failure(primary) - - assert ctx.terminal_outcome == "error" - assert not ctx.evidence_complete - - @pytest.mark.asyncio async def test_terminal_evidence_emit_failure_still_closes_sink(): sink = RecordingSink(fail_emit=RunEventKind.RUN_SUCCEEDED) @@ -541,90 +268,10 @@ def mount_file_at(self, host_path: str, sandbox_path: str) -> None: def mkdir_p(self, sandbox_path: str) -> None: self._path(sandbox_path).mkdir(parents=True, exist_ok=True) - def list_dir(self, sandbox_path: str) -> list[str]: - root = self._path(sandbox_path) - return [ - f"{sandbox_path.rstrip('/')}/{path.relative_to(root).as_posix()}" - for path in sorted(root.rglob("*")) - if path.is_file() - ] - - def sync_file_to(self, sandbox_path: str, host_path: str) -> None: - target = Path(host_path) - target.parent.mkdir(parents=True, exist_ok=True) - shutil.copy2(self._path(sandbox_path), target) - def shutdown(self) -> None: self.shutdown_calls += 1 -@pytest.mark.asyncio -async def test_interpreter_session_round_trips_directory_artifacts(tmp_path: Path): - source = tmp_path / "source" - source.mkdir() - (source / "nested").mkdir() - (source / "nested" / "value.txt").write_text("value", encoding="utf-8") - interpreter = FileInterpreter(tmp_path / "sandbox") - session = InterpreterExecutionSession( - interpreter, - name="test", - ownership=SessionOwnership.INJECTED, - ) - input_artifact = Artifact( - id="directory-input", - kind="opaque", - metadata={ - "source_path": str(source), - "sandbox_path": "/sandbox/input/source", - }, - ) - - binding = await session.mount(input_artifact) - output_dir = tmp_path / "collected" - collected = await session.collect( - Artifact( - id="directory-output", - kind="opaque", - metadata={ - "sandbox_path": binding.path, - "destination_path": str(output_dir), - "directory": True, - }, - ) - ) - - assert collected == str(output_dir) - assert (output_dir / "nested" / "value.txt").read_text(encoding="utf-8") == "value" - - -@pytest.mark.asyncio -async def test_interpreter_backend_adapter_preserves_injected_ownership(tmp_path: Path): - interpreter = FileInterpreter(tmp_path / "sandbox") - exits = 0 - - @contextmanager - def acquire(spec: ExecutionSpec, ctx: RunContext): - nonlocal exits - yield interpreter - exits += 1 - - backend = InterpreterBackendAdapter( - "injected", - acquire, - ownership=SessionOwnership.INJECTED, - supports_host_directory_mounts=True, - ) - spec = make_spec() - ctx = RunContext(spec, {}) - - async with backend.start(ExecutionSpec(), ctx) as session: - result = await session.run_code("code", {"value": 3}) - assert result.value == "code:3" - - assert exits == 1 - assert interpreter.shutdown_calls == 0 - - @pytest.mark.asyncio async def test_injected_backend_cancellation_does_not_leak_waiting_lock(tmp_path: Path): interpreter = FileInterpreter(tmp_path / "sandbox") @@ -642,9 +289,9 @@ def acquire(spec: ExecutionSpec, ctx: RunContext): first_ctx = RunContext(make_spec(), {}) waiting_ctx = RunContext(make_spec(), {}) - async def wait_for_session() -> None: - async with backend.start(ExecutionSpec(), waiting_ctx): - pass + async def wait_for_session() -> str: + async with backend.start(ExecutionSpec(), waiting_ctx) as session: + return (await session.run_code("code", {"value": 3})).value async with backend.start(ExecutionSpec(), first_ctx): waiting = asyncio.create_task(wait_for_session()) @@ -653,31 +300,8 @@ async def wait_for_session() -> None: with pytest.raises(asyncio.CancelledError): await waiting - await asyncio.sleep(0.05) - assert backend._invocation_lock is not None - assert backend._invocation_lock.acquire(blocking=False) - backend._invocation_lock.release() - - -@pytest.mark.asyncio -async def test_injected_backend_releases_lock_when_acquisition_raises(): - def acquire(spec: ExecutionSpec, ctx: RunContext): - raise RuntimeError("acquire failed") - - backend = InterpreterBackendAdapter( - "injected", - acquire, - ownership=SessionOwnership.INJECTED, - supports_host_directory_mounts=True, - ) - - with pytest.raises(RuntimeError, match="acquire failed"): - async with backend.start(ExecutionSpec(), RunContext(make_spec(), {})): - pass - - assert backend._invocation_lock is not None - assert backend._invocation_lock.acquire(blocking=False) - backend._invocation_lock.release() + assert await asyncio.wait_for(wait_for_session(), timeout=2) == "code:3" + assert interpreter.shutdown_calls == 0 @pytest.mark.asyncio @@ -717,9 +341,7 @@ def acquire(spec: ExecutionSpec, ctx: RunContext): second_ctx = RunContext(make_spec(), {}) second_mount = HostDirectoryMount(str(second), "/workspace") - async with backend.start( - ExecutionSpec(host_directory_mounts=(second_mount,)), second_ctx - ): + async with backend.start(ExecutionSpec(host_directory_mounts=(second_mount,)), second_ctx): pass assert attempts == 2 @@ -744,71 +366,48 @@ def acquire(spec, ctx): mount = HostDirectoryMount(str(tmp_path), "/workspace") with pytest.raises(RuntimeError, match="pooled interpreters"): - async with backend.start( - ExecutionSpec(host_directory_mounts=(mount,)), ctx - ): + async with backend.start(ExecutionSpec(host_directory_mounts=(mount,)), ctx): pass assert acquisitions == 0 @pytest.mark.asyncio -async def test_unsupported_injected_mount_is_rejected_before_acquisition(tmp_path: Path): - acquisitions = 0 - - @contextmanager - def acquire(spec, ctx): - nonlocal acquisitions - acquisitions += 1 - yield FileInterpreter(tmp_path / "sandbox") - - backend = InterpreterBackendAdapter( - "injected", - acquire, +async def test_session_rejects_undeclared_host_directory_mount(tmp_path: Path): + session = InterpreterExecutionSession( + FileInterpreter(tmp_path / "sandbox"), + name="injected", ownership=SessionOwnership.INJECTED, ) - mount = HostDirectoryMount(str(tmp_path), "/workspace") - with pytest.raises(RuntimeError, match="does not support host directory mounts"): - async with backend.start( - ExecutionSpec(host_directory_mounts=(mount,)), - RunContext(make_spec(), {}), - ): - pass - - assert acquisitions == 0 + with pytest.raises(RuntimeError, match="declared before backend acquisition"): + await session.mount_host_directory(HostDirectoryMount(str(tmp_path), "/workspace")) @pytest.mark.asyncio -async def test_backend_without_mount_capability_rejects_before_start(tmp_path: Path): - from predict_rlm.backends.adapters import ExistingExecutionBackendAdapter - - backend = FinalBackend() - wrapped = ExistingExecutionBackendAdapter(backend) - - with pytest.raises(RuntimeError, match="does not support host directory mounts"): - await wrapped.validate_host_directory_mounts( - (HostDirectoryMount(str(tmp_path), "/workspace"),), - RunContext(make_spec(), {}), - ) - - assert backend.spec is None +async def test_injected_backend_waits_for_cancelled_sync_execution_before_release(): + class BlockingInterpreter: + def __init__(self) -> None: + self.started = threading.Event() + self.release = threading.Event() + self.finished = threading.Event() + def execute(self, code, variables=None, timeout=None): + self.started.set() + self.release.wait() + self.finished.set() + return "finished" -@pytest.mark.asyncio -@pytest.mark.parametrize("change", ["add", "removal", "access"]) -async def test_injected_backend_rejects_semantic_mount_aggregate_changes( - tmp_path: Path, - change: str, -): - acquisitions = 0 - interpreter = FileInterpreter(tmp_path / "sandbox") + interpreter = BlockingInterpreter() + finished_on_exit = False @contextmanager - def acquire(spec, ctx): - nonlocal acquisitions - acquisitions += 1 - yield interpreter + def acquire(spec: ExecutionSpec, ctx: RunContext): + nonlocal finished_on_exit + try: + yield interpreter + finally: + finished_on_exit = interpreter.finished.is_set() backend = InterpreterBackendAdapter( "injected", @@ -816,377 +415,32 @@ def acquire(spec, ctx): ownership=SessionOwnership.INJECTED, supports_host_directory_mounts=True, ) - first = HostDirectoryMount(str(tmp_path / "first"), "/first") - second = HostDirectoryMount(str(tmp_path / "second"), "/second") - if change == "add": - initial = () - changed = (first,) - elif change == "removal": - initial = (first, second) - changed = (first,) - else: - initial = (first,) - changed = ( - HostDirectoryMount(first.host_path, first.sandbox_path, read_only=True), - ) - - async with backend.start( - ExecutionSpec(host_directory_mounts=initial), - RunContext(make_spec(), {}), - ): - pass + async with backend.start(ExecutionSpec(), RunContext(make_spec(), {})) as session: + execution = asyncio.create_task(session.run_code("block")) + await asyncio.to_thread(interpreter.started.wait) + execution.cancel() + await asyncio.sleep(0.05) + assert not execution.done() - with pytest.raises(ValueError, match="mount set"): - async with backend.start( - ExecutionSpec(host_directory_mounts=changed), - RunContext(make_spec(), {}), - ): - pass + interpreter.release.set() + with pytest.raises(asyncio.CancelledError): + await execution - assert acquisitions == 1 + assert finished_on_exit + assert interpreter.finished.is_set() @pytest.mark.asyncio -async def test_injected_backend_accepts_semantically_reordered_mount_aggregate( - tmp_path: Path, -): - acquisitions = 0 +async def test_cancelled_execution_defers_adapter_work_until_session_is_idle(): + from predict_rlm import PredictRLM - @contextmanager - def acquire(spec, ctx): - nonlocal acquisitions - acquisitions += 1 - yield FileInterpreter(tmp_path / "sandbox") + started = asyncio.Event() + after_calls = [] + finalize_live_states = [] - backend = InterpreterBackendAdapter( - "injected", - acquire, - ownership=SessionOwnership.INJECTED, - supports_host_directory_mounts=True, - ) - first = HostDirectoryMount(str(tmp_path / "first"), "/first") - second = HostDirectoryMount(str(tmp_path / "second"), "/second") - - async with backend.start( - ExecutionSpec(host_directory_mounts=(first, second)), - RunContext(make_spec(), {}), - ): - pass - async with backend.start( - ExecutionSpec(host_directory_mounts=(second, first)), - RunContext(make_spec(), {}), - ): - pass - - assert acquisitions == 2 - - -@pytest.mark.asyncio -async def test_async_only_injected_mount_configuration_is_rejected_before_acquisition( - tmp_path: Path, -): - from predict_rlm.compatibility.backends import execution_from_options - - acquisitions = 0 - - class AsyncOnlyInterpreter: - async def aconfigure_direct_workspace_mounts(self, mounts): - raise AssertionError("sync compatibility adapter cannot call this hook") - - @contextmanager - def acquire(spec, ctx): - nonlocal acquisitions - acquisitions += 1 - yield AsyncOnlyInterpreter() - - owner = MagicMock() - owner._acquire_runtime_interpreter = acquire - backend = execution_from_options( - owner=owner, - interpreter=AsyncOnlyInterpreter(), - sandbox_backend=MagicMock(value="injected"), - sbx_config=None, - sbx_pool=None, - allowed_domains=None, - runtime_hooks=[], - on_runtime_hook_event=None, - ) - mount = HostDirectoryMount(str(tmp_path), "/workspace") - - with pytest.raises(RuntimeError, match="does not support host directory mounts"): - async with backend.start( - ExecutionSpec(host_directory_mounts=(mount,)), - RunContext(make_spec(), {}), - ): - pass - - assert acquisitions == 0 - - -@pytest.mark.asyncio -async def test_session_rejects_undeclared_host_directory_mount(tmp_path: Path): - session = InterpreterExecutionSession( - FileInterpreter(tmp_path / "sandbox"), - name="injected", - ownership=SessionOwnership.INJECTED, - ) - - with pytest.raises(RuntimeError, match="declared before backend acquisition"): - await session.mount_host_directory( - HostDirectoryMount(str(tmp_path), "/workspace") - ) - - -@pytest.mark.asyncio -async def test_backend_acquisition_receives_body_exception(tmp_path: Path): - interpreter = FileInterpreter(tmp_path / "sandbox") - exit_args = None - - class Manager: - def __enter__(self): - return interpreter - - def __exit__(self, exc_type, exc, traceback): - nonlocal exit_args - exit_args = (exc_type, exc, traceback) - - backend = InterpreterBackendAdapter( - "owned", - lambda spec, ctx: Manager(), - ownership=SessionOwnership.OWNED, - ) - - with pytest.raises(ValueError, match="body failed"): - async with backend.start(ExecutionSpec(), RunContext(make_spec(), {})): - raise ValueError("body failed") - - assert exit_args is not None - assert exit_args[0] is ValueError - assert str(exit_args[1]) == "body failed" - - -@pytest.mark.asyncio -async def test_injected_backend_waits_for_cancelled_sync_execution_before_release(): - class BlockingInterpreter: - def __init__(self) -> None: - self.started = threading.Event() - self.release = threading.Event() - self.finished = threading.Event() - - def execute(self, code, variables=None, timeout=None): - self.started.set() - self.release.wait() - self.finished.set() - return "finished" - - interpreter = BlockingInterpreter() - finished_on_exit = False - - @contextmanager - def acquire(spec: ExecutionSpec, ctx: RunContext): - nonlocal finished_on_exit - try: - yield interpreter - finally: - finished_on_exit = interpreter.finished.is_set() - - backend = InterpreterBackendAdapter( - "injected", - acquire, - ownership=SessionOwnership.INJECTED, - supports_host_directory_mounts=True, - ) - async with backend.start(ExecutionSpec(), RunContext(make_spec(), {})) as session: - execution = asyncio.create_task(session.run_code("block")) - await asyncio.to_thread(interpreter.started.wait) - execution.cancel() - await asyncio.sleep(0.05) - assert not execution.done() - - interpreter.release.set() - with pytest.raises(asyncio.CancelledError): - await execution - - assert finished_on_exit - assert interpreter.finished.is_set() - - -@pytest.mark.asyncio -async def test_jspi_cancellation_cancels_pending_host_calls(): - from predict_rlm.backends.jspi import JspiBackend - - started = asyncio.Event() - cancelled = asyncio.Event() - child_task = None - - class FakeJspi: - _execute_async_loop = JspiBackend._execute_async_loop - - def __init__(self) -> None: - self._active_tool_count = 0 - self._pending_file_ops = {} - self.reads = 0 - - async def _send_completed_responses(self, pending_tasks) -> None: - return None - - async def _read_with_timeout_async(self, timeout): - self.reads += 1 - if self.reads == 1: - return ( - '{"jsonrpc":"2.0","id":"tool-1","method":"tool_call",' - '"params":{"name":"slow","args":[],"kwargs":{}}}' - ) - await asyncio.Event().wait() - - async def _execute_tool_async(self, name, params, request_id): - nonlocal child_task - child_task = asyncio.current_task() - started.set() - try: - await asyncio.Event().wait() - finally: - cancelled.set() - - fake = FakeJspi() - execution = asyncio.create_task(JspiBackend._execute_async(fake, 1)) - await started.wait() - execution.cancel() - - with pytest.raises(asyncio.CancelledError): - await execution - - try: - assert cancelled.is_set() - assert fake._active_tool_count == 0 - finally: - if child_task is not None and not child_task.done(): - child_task.cancel() - await asyncio.gather(child_task, return_exceptions=True) - - -def test_backends_execution_backend_exports_final_session_contract(): - from predict_rlm.backends import ExecutionBackend - - assert hasattr(ExecutionBackend, "start") - assert not hasattr(ExecutionBackend, "execute") - - -def test_maintained_backends_select_native_session_adapter(): - from predict_rlm import PredictRLM - from predict_rlm.backends.jspi import JspiExecutionBackend - - lm = MagicMock() - lm.copy.return_value = lm - lm.history = [] - - rlm = PredictRLM("question -> answer", lm=lm) - - assert isinstance(rlm.runtime_spec.execution, JspiExecutionBackend) - - -@pytest.mark.sbx -def test_maintained_sbx_backends_select_final_execution_ownership_seam(): - from predict_rlm import PredictRLM - from predict_rlm.backends.sbx import ( - SbxExecutionBackend, - SbxPool, - SbxPoolExecutionBackend, - ) - - lm = MagicMock() - lm.copy.return_value = lm - lm.history = [] - - owned = PredictRLM("question -> answer", lm=lm, sandbox_backend="sbx") - pool = SbxPool(size=1, preinstall_packages=False) - pooled = PredictRLM( - "question -> answer", - lm=lm, - sandbox_backend="sbx", - sbx_pool=pool, - ) - - assert isinstance(owned.runtime_spec.execution, SbxExecutionBackend) - assert isinstance(pooled.runtime_spec.execution, SbxPoolExecutionBackend) - - -@pytest.mark.asyncio -async def test_input_adapter_after_execution_covers_generated_success_and_failure_only( - tmp_path: Path, -): - from predict_rlm import PredictRLM - - completed = [] - - class LifecycleInputAdapter(InputAdapter[str]): - name = "lifecycle" - value_type = str - - async def prepare(self, field, value, ctx): - return PreparedInput(model_value=value) - - async def after_execution( - self, - field, - prepared, - ctx, - session, - result, - error, - ): - completed.append((field.name, result, error)) - - class FailingRunSession(FinalSession): - async def run_code(self, code, variables=None, timeout=None): - if code == "raise RuntimeError": - raise RuntimeError("execution failed") - return await super().run_code(code, variables, timeout=timeout) - - backend = FinalBackend() - backend.session = FailingRunSession() - rlm = PredictRLM( - "value: str -> answer: str", - lm=MagicMock(history=[]), - execution=backend, - adapters=[LifecycleInputAdapter()], - ) - module_path = tmp_path / "helper.py" - module_path.write_text("VALUE = 1\n", encoding="utf-8") - rlm._skill_modules = {"helper": str(module_path)} - ctx = rlm._new_run_context({"value": "input"}) - - async with use_run_context(ctx): - await rlm._prepare_runtime_inputs(ctx, ctx.input_values) - async with rlm._execution_session({}) as repl: - await rlm._bind_runtime_inputs(ctx) - await repl.aexecute("success") - with pytest.raises(RuntimeError, match="execution failed"): - await repl.aexecute("raise RuntimeError") - await rlm._setup_runtime_modules(ctx) - - assert len(completed) == 2 - outcomes = [ - (result is not None, type(error) if error else None) - for _, result, error in completed - ] - assert outcomes == [ - (True, None), - (False, RuntimeError), - ] - - -@pytest.mark.asyncio -async def test_cancelled_execution_defers_adapter_work_until_session_is_idle(): - from predict_rlm import PredictRLM - - started = asyncio.Event() - after_calls = [] - finalize_live_states = [] - - class LifecycleInputAdapter(InputAdapter[str]): - name = "cancel-lifecycle" - value_type = str + class LifecycleInputAdapter(InputAdapter[str]): + name = "cancel-lifecycle" + value_type = str async def prepare(self, field, value, ctx): return PreparedInput(model_value=value) @@ -1226,17 +480,14 @@ async def cancel(self): lm=MagicMock(history=[]), execution=backend, adapters=[LifecycleInputAdapter()], + max_iterations=1, + verbose=False, + ) + rlm.generate_action.acall = AsyncMock( + return_value=dspy.Prediction(reasoning="block", code="await block()") ) - ctx = rlm._new_run_context({"value": "input"}) - - async def invoke(): - async with use_run_context(ctx): - await rlm._prepare_runtime_inputs(ctx, ctx.input_values) - async with rlm._execution_session({}) as repl: - await rlm._bind_runtime_inputs(ctx) - await repl.aexecute("block") - invocation = asyncio.create_task(invoke()) + invocation = asyncio.create_task(rlm.aforward(value="input")) await started.wait() invocation.cancel() with pytest.raises(asyncio.CancelledError): @@ -1360,58 +611,6 @@ async def finalize(self): assert isinstance(getattr(raised.value, attribute), OSError) -@pytest.mark.asyncio -async def test_input_adapter_finalization_is_reverse_and_idempotent(): - from predict_rlm import PredictRLM - - calls = [] - - class StringInputAdapter(InputAdapter[str]): - name = "string-lifecycle" - value_type = str - - async def prepare(self, field, value, ctx): - return PreparedInput(model_value=value) - - async def finalize(self, field, prepared, ctx, session, error): - calls.append(field.name) - - class IntegerInputAdapter(InputAdapter[int]): - name = "integer-lifecycle" - value_type = int - - async def prepare(self, field, value, ctx): - return PreparedInput(model_value=value) - - async def finalize(self, field, prepared, ctx, session, error): - calls.append(field.name) - - class RecordingFinalizeSession(FinalSession): - async def finalize(self): - calls.append("session") - await super().finalize() - - backend = FinalBackend() - backend.session = RecordingFinalizeSession() - rlm = PredictRLM( - "first: str, second: int -> answer: str", - lm=MagicMock(history=[]), - execution=backend, - adapters=[StringInputAdapter(), IntegerInputAdapter()], - ) - ctx = rlm._new_run_context({"first": "one", "second": 2}) - ctx.session = backend.session - - async with use_run_context(ctx): - await rlm._prepare_runtime_inputs(ctx, ctx.input_values) - await rlm._open_runtime_inputs(ctx) - await rlm._bind_runtime_inputs(ctx) - await rlm._finalize_runtime_inputs(ctx, backend.session, None) - await rlm._finalize_runtime_inputs(ctx, backend.session, None) - - assert calls == ["second", "first", "session"] - - class FinalSession: name = "final" ownership = SessionOwnership.OWNED @@ -1419,15 +618,11 @@ class FinalSession: def __init__(self) -> None: self.finalized = 0 self.cancelled = 0 - self.mounted = [] - self.variables = None - self.final_payload = None async def install_packages(self, packages) -> None: return None async def mount(self, artifact): - self.mounted.append(artifact) return ArtifactBinding( artifact_id=artifact.id, path=artifact.metadata["sandbox_path"], @@ -1436,10 +631,7 @@ async def mount(self, artifact): async def run_code(self, code, variables=None, timeout=None): from predict_rlm.runtime import ExecutionResult - self.variables = variables - payload = self.final_payload or { - "answer": (variables or {}).get("question", "async-path") - } + payload = {"answer": (variables or {}).get("question", "async-path")} return ExecutionResult(FinalOutput(payload)) async def collect(self, artifact): @@ -1484,116 +676,20 @@ async def start(self, spec, ctx): @pytest.mark.asyncio -async def test_external_input_adapter_owns_session_lifecycle(): +async def test_session_finalizes_after_input_adapter_failure_in_reverse_order(): from predict_rlm import PredictRLM - from predict_rlm.runtime import HostDirectoryMount calls = [] - host_mount = HostDirectoryMount("/host/external", "/external") - class LifecycleInputAdapter(InputAdapter[str]): - name = "lifecycle" + class StringInputAdapter(InputAdapter[str]): + name = "string-finalize" value_type = str async def prepare(self, field, value, ctx): - calls.append("prepare") - return PreparedInput( - model_value="planned", - artifacts=(Artifact(id="external", kind="external"),), - host_directory_mounts=(host_mount,), - ) - - async def open(self, field, prepared, ctx, backend): - calls.append(("open", field.name, backend.name)) - - async def bind(self, field, prepared, ctx, session): - calls.append("bind") - path = await session.mount_host_directory(host_mount) - return BoundInput( - model_value="bound", - bindings=(ArtifactBinding(artifact_id="external", path=path),), - ) - - async def after_execution( - self, - field, - prepared, - ctx, - session, - result, - error, - ): - calls.append(("after_execution", field.name, result.value, error)) + return PreparedInput(model_value=value) async def finalize(self, field, prepared, ctx, session, error): - calls.append(("finalize", field.name, error)) - - class RecordingBackend(FinalBackend): - async def validate_host_directory_mounts(self, mounts, ctx): - calls.append(("validate_host_directory_mounts", tuple(mounts))) - - @asynccontextmanager - async def start(self, spec, ctx): - calls.append("start") - async with super().start(spec, ctx) as session: - yield session - calls.append("release") - - class RecordingHostMountSession(FinalSession): - def __init__(self): - super().__init__() - self.host_directory_mounts = [] - - async def mount_host_directory(self, mount): - self.host_directory_mounts.append(mount) - return mount.sandbox_path - - backend = RecordingBackend() - backend.session = RecordingHostMountSession() - rlm = PredictRLM( - "value: str -> answer: str", - lm=MagicMock(history=[]), - execution=backend, - adapters=[LifecycleInputAdapter()], - ) - ctx = rlm._new_run_context({"value": "input"}) - - async with use_run_context(ctx): - await rlm._prepare_runtime_inputs(ctx, ctx.input_values) - async with rlm._execution_session({}) as repl: - await rlm._bind_runtime_inputs(ctx) - assert ctx.input_bindings["value"].prepared.model_value == "bound" - await repl.aexecute("print('run')") - - assert calls[:5] == [ - "prepare", - ("open", "value", "final"), - ("validate_host_directory_mounts", (host_mount,)), - "start", - "bind", - ] - assert calls[5][0] == "after_execution" - assert calls[6:] == [("finalize", "value", None), "release"] - assert backend.session.host_directory_mounts == [host_mount] - assert backend.session.mounted == [] - assert backend.session.finalized == 1 - - -@pytest.mark.asyncio -async def test_session_finalizes_after_input_adapter_failure_in_reverse_order(): - from predict_rlm import PredictRLM - - calls = [] - - class StringInputAdapter(InputAdapter[str]): - name = "string-finalize" - value_type = str - - async def prepare(self, field, value, ctx): - return PreparedInput(model_value=value) - - async def finalize(self, field, prepared, ctx, session, error): - calls.append(field.name) + calls.append(field.name) class IntegerInputAdapter(InputAdapter[int]): name = "integer-finalize" @@ -1636,8 +732,6 @@ async def finalize(self): ("failure_stage", "expected_fields", "session_available"), [ ("open", ["second", "first"], False), - ("acquisition", ["third", "second", "first"], False), - ("install", ["third", "second", "first"], True), ("bind", ["third", "second", "first"], True), ], ) @@ -1671,38 +765,18 @@ async def bind(self, field, prepared, ctx, session): async def finalize(self, field, prepared, ctx, session, error): finalized.append((field.name, session is not None)) - class StartupSession(FinalSession): - async def install_packages(self, packages): - if failure_stage == "install": - raise RuntimeError("install failed") - - class StartupBackend(FinalBackend): - @asynccontextmanager - async def start(self, spec, ctx): - if failure_stage == "acquisition": - raise RuntimeError("acquisition failed") - async with super().start(spec, ctx) as session: - yield session - - backend = StartupBackend() - backend.session = StartupSession() + backend = FinalBackend() rlm = PredictRLM( "first: str, second: int, third: bool -> answer: str", lm=MagicMock(history=[]), execution=backend, adapters=[LifecycleAdapter()], ) - ctx = rlm._new_run_context({"first": "one", "second": 2, "third": True}) - with pytest.raises(RuntimeError, match=failure_stage.replace("_", " ")): - async with use_run_context(ctx): - await rlm._prepare_runtime_inputs(ctx, ctx.input_values) - async with rlm._execution_session({}): - await rlm._bind_runtime_inputs(ctx) + with pytest.raises(RuntimeError, match=failure_stage): + await rlm.aforward(first="one", second=2, third=True) - assert finalized == [ - (field_name, session_available) for field_name in expected_fields - ] + assert finalized == [(field_name, session_available) for field_name in expected_fields] if failure_stage == "open": assert entered == ["first", "second"] @@ -1747,91 +821,6 @@ async def finalize(self, field, prepared, ctx, session, error): } -@pytest.mark.asyncio -async def test_input_sandbox_root_collision_is_rejected_before_acquisition(): - from predict_rlm import PredictRLM, SandboxRootReservation - - acquisitions = 0 - - class FirstAdapter(InputAdapter[str]): - name = "first-root" - value_type = str - - async def prepare(self, field, value, ctx): - return PreparedInput( - model_value=value, - sandbox_roots=(SandboxRootReservation("/repository"),), - ) - - class SecondAdapter(InputAdapter[int]): - name = "second-root" - value_type = int - - async def prepare(self, field, value, ctx): - return PreparedInput( - model_value=value, - sandbox_roots=(SandboxRootReservation("/repository/subdir"),), - ) - - class AcquisitionBackend(FinalBackend): - @asynccontextmanager - async def start(self, spec, ctx): - nonlocal acquisitions - acquisitions += 1 - async with super().start(spec, ctx) as session: - yield session - - rlm = PredictRLM( - "first: str, second: int -> answer: str", - lm=MagicMock(history=[]), - execution=AcquisitionBackend(), - adapters=[FirstAdapter(), SecondAdapter()], - ) - ctx = rlm._new_run_context({"first": "one", "second": 2}) - - with pytest.raises(ValueError, match="sandbox destination.*overlap"): - async with use_run_context(ctx): - await rlm._prepare_runtime_inputs(ctx, ctx.input_values) - async with rlm._execution_session({}): - pass - - assert acquisitions == 0 - - -@pytest.mark.asyncio -async def test_output_reservation_rejects_prepared_input_overlap(tmp_path: Path): - from predict_rlm import PredictRLM - - source = tmp_path / "source.txt" - source.write_text("source", encoding="utf-8") - - class CollidingFileInputAdapter(FileInputAdapter): - async def prepare(self, field, value, ctx): - item = field.unpack(value)[0] - assert isinstance(item, RuntimeFile) - assert item.path is not None - return PreparedInput.path(item.path, at="output/result") - - rlm = PredictRLM( - KernelFileSignature, - lm=MagicMock(history=[]), - execution=FinalBackend(), - adapters=[CollidingFileInputAdapter()], - ) - ctx = rlm._new_run_context({"source": RuntimeFile(path=str(source))}) - - async with use_run_context(ctx): - await rlm._prepare_runtime_inputs(ctx, dict(ctx.input_values)) - await rlm._prepare_runtime_outputs(ctx) - async with rlm._execution_session({}): - await rlm._bind_runtime_inputs(ctx) - with pytest.raises( - ValueError, - match="Input/output sandbox destinations overlap", - ): - await rlm._reserve_runtime_outputs(ctx) - - @pytest.mark.asyncio async def test_transfer_root_host_mount_collision_is_rejected_before_acquisition( tmp_path: Path, @@ -1857,9 +846,7 @@ class MountAdapter(InputAdapter[int]): async def prepare(self, field, value, ctx): return PreparedInput( model_value=value, - host_directory_mounts=( - HostDirectoryMount(str(tmp_path), "/repository"), - ), + host_directory_mounts=(HostDirectoryMount(str(tmp_path), "/repository"),), ) class AcquisitionBackend(FinalBackend): @@ -1907,9 +894,7 @@ async def test_interpreter_transfer_operations_require_declared_sandbox_root( ) await session.create_directory("/repository") - await session.transfer_file( - FileTransfer(str(source), "/repository/source.txt") - ) + await session.transfer_file(FileTransfer(str(source), "/repository/source.txt")) with pytest.raises(UnsupportedOperationError, match="declared sandbox roots"): await session.create_directory("/other") @@ -1917,41 +902,6 @@ async def test_interpreter_transfer_operations_require_declared_sandbox_root( await session.transfer_file(FileTransfer(str(source), "/other/source.txt")) -@pytest.mark.asyncio -async def test_input_adapter_finalize_failure_preserves_primary_and_evidence(): - from predict_rlm import PredictRLM - - class FailingInputAdapter(InputAdapter[str]): - name = "failing-finalize" - value_type = str - - async def prepare(self, field, value, ctx): - return PreparedInput(model_value=value) - - async def finalize(self, field, prepared, ctx, session, error): - raise OSError("adapter finalize failed") - - backend = FinalBackend() - rlm = PredictRLM( - "question: str -> answer: str", - lm=MagicMock(history=[]), - execution=backend, - adapters=[FailingInputAdapter()], - max_iterations=1, - verbose=False, - ) - rlm.generate_action.acall = AsyncMock(side_effect=ValueError("primary failure")) - - with pytest.raises(ValueError, match="primary failure") as raised: - await rlm.aforward(question="test") - - assert isinstance(raised.value.input_adapter_finalize_error, OSError) - assert raised.value.trace.evidence.complete is False - assert "session.finalize_failed" in { - event.kind for event in raised.value.trace.evidence.events - } - - class BlockingFinalizeSession(FinalSession): def __init__(self) -> None: super().__init__() @@ -1992,341 +942,32 @@ def __init__(self) -> None: class SyncedFinalSession(FinalSession): - def __init__(self) -> None: - super().__init__() - self.files = {"/sandbox/work.txt": "before"} - self.synced_mounts = [] - self.tool_name = "mutate" - self.captured_host_path = None - self.spec = None - async def run_code(self, code, variables=None, timeout=None): - tool = self.spec.tools[self.tool_name] + tool = self.spec.tools["block"] await tool("/sandbox/work.txt") return await super().run_code(code, variables, timeout=timeout) async def collect(self, artifact): destination = Path(artifact.metadata["destination_path"]) destination.parent.mkdir(parents=True, exist_ok=True) - destination.write_text( - self.files[artifact.metadata["sandbox_path"]], - encoding="utf-8", - ) + destination.write_text("before", encoding="utf-8") return str(destination) - async def mount(self, artifact): - if artifact.kind == "compat.file" and "destination_path" not in artifact.metadata: - source = Path(artifact.metadata["source_path"]) - self.files[artifact.metadata["sandbox_path"]] = source.read_text( - encoding="utf-8" - ) - self.synced_mounts.append(artifact) - return ArtifactBinding( - artifact_id=artifact.id, - path=artifact.metadata["sandbox_path"], - ) - - -class SyncedFinalBackend(FinalBackend): - def __init__(self) -> None: - self.session = SyncedFinalSession() - self.spec = None - - -class MaintainedSyncedInterpreter: - def __init__(self) -> None: - self.files = {"/sandbox/work.txt": "before"} - self.tools = {} - self.shutdown_calls = 0 - - async def aensure_skill_packages(self, packages) -> None: - return None - - async def aexecute(self, code, variables=None, timeout=None): - await self.tools["mutate"]("/sandbox/work.txt") - return FinalOutput({"answer": "done"}) - - async def async_file_to(self, sandbox_path: str, host_path: str) -> None: - destination = Path(host_path) - destination.parent.mkdir(parents=True, exist_ok=True) - destination.write_text(self.files[sandbox_path], encoding="utf-8") - - async def amount_file_at(self, host_path: str, sandbox_path: str) -> None: - self.files[sandbox_path] = Path(host_path).read_text(encoding="utf-8") - - async def _sync_file_during_tool(self, sandbox_path: str, host_path: str) -> None: - await self.async_file_to(sandbox_path, host_path) - - async def _mount_file_during_tool(self, host_path: str, sandbox_path: str) -> None: - await self.amount_file_at(host_path, sandbox_path) - - async def ainterrupt(self) -> None: - return None - - def retire_when_sync_workers_finish(self) -> bool: - return False - - def retire_when_host_work_finishes(self) -> bool: - return False - - async def ashutdown(self) -> None: - self.shutdown_calls += 1 - - -class MaintainedSyncedPool: - def __init__(self, interpreter: MaintainedSyncedInterpreter) -> None: - self._interpreter_kwargs = {} - self.session_requirements = SessionRequirements() - self.interpreter = interpreter - self.released = False - - @asynccontextmanager - async def alease(self, **kwargs): - self.interpreter.tools = kwargs["tools"] - try: - yield self.interpreter - finally: - self.released = True - - -def test_predict_rlm_sync_entry_uses_async_session_contract(): - from predict_rlm import PredictRLM - - lm = MagicMock() - lm.copy.return_value = lm - lm.history = [] - backend = FinalBackend() - rlm = PredictRLM( - "question -> answer", - lm=lm, - execution=backend, - max_iterations=1, - verbose=False, - ) - rlm.generate_action.acall = AsyncMock( - return_value=dspy.Prediction( - reasoning="submit", - code="SUBMIT(answer='async-path')", - ) - ) - rlm._build_signatures_with_files = MagicMock( - return_value=(rlm.generate_action, rlm.extract) - ) - - result = rlm.forward(question="test") - - assert result.answer == "test" - assert result.trace.status == "completed" - assert backend.session.finalized >= 1 - - -@pytest.mark.asyncio -async def test_failed_backend_exit_emits_release_failure_not_release_success(): - from predict_rlm import PredictRLM - - sink = RecordingSink() - backend = FailingExitBackend() - rlm = PredictRLM( - "question: str -> answer: str", - lm=MagicMock(history=[]), - execution=backend, - events=(sink,), - max_iterations=1, - verbose=False, - ) - rlm.generate_action.acall = AsyncMock( - return_value=dspy.Prediction( - reasoning="submit", - code="SUBMIT(answer=question)", - ) - ) - - with pytest.raises(OSError, match="release failed"): - await rlm.aforward(question="test") - - kinds = [event.kind for event in sink.events] - assert RunEventKind.SESSION_RELEASED not in kinds - assert RunEventKind.SESSION_RELEASE_FAILED in kinds - - -@pytest.mark.asyncio -async def test_failed_backend_exit_does_not_mask_primary_execution_error(): - from predict_rlm import PredictRLM - - backend = FailingExitBackend() - rlm = PredictRLM( - "question: str -> answer: str", - lm=MagicMock(history=[]), - execution=backend, - max_iterations=1, - verbose=False, - ) - rlm.generate_action.acall = AsyncMock( - side_effect=ValueError("primary execution failure") - ) - - with pytest.raises(ValueError, match="primary execution failure") as raised: - await rlm.aforward(question="test") - - assert isinstance(raised.value.session_release_error, OSError) - - -@pytest.mark.asyncio -async def test_cancellation_waits_for_owned_finalization_before_backend_release(): - from predict_rlm import PredictRLM - - backend = BlockingFinalizeBackend() - rlm = PredictRLM( - "question: str -> answer: str", - lm=MagicMock(history=[]), - execution=backend, - max_iterations=1, - verbose=False, - ) - rlm.generate_action.acall = AsyncMock( - return_value=dspy.Prediction( - reasoning="submit", - code="SUBMIT(answer=question)", - ) - ) - - invocation = asyncio.create_task(rlm.aforward(question="test")) - await backend.session.finalize_started.wait() - invocation.cancel() - await asyncio.sleep(0) - completed_before_finalize = invocation.done() - released_before_finalize = backend.released - backend.session.allow_finalize.set() - with pytest.raises(asyncio.CancelledError): - await invocation - - assert not completed_before_finalize - assert not released_before_finalize - assert backend.session.finalized == 1 - assert backend.released - - -@pytest.mark.asyncio -async def test_finalization_failure_preserves_primary_and_marks_evidence_incomplete(): - from predict_rlm import PredictRLM - - backend = FailingFinalizeBackend() - rlm = PredictRLM( - "question: str -> answer: str", - lm=MagicMock(history=[]), - execution=backend, - max_iterations=1, - verbose=False, - ) - rlm.generate_action.acall = AsyncMock(side_effect=ValueError("primary failure")) - - with pytest.raises(ValueError, match="primary failure") as raised: - await rlm.aforward(question="test") - - assert isinstance(raised.value.session_finalize_error, OSError) - assert raised.value.trace.evidence.complete is False - assert "session.finalize_failed" in { - event.kind for event in raised.value.trace.evidence.events - } - - -def test_run_trace_strict_evidence_reaches_rlm_gepa_consumer(): - from predict_rlm import PredictRLM - from rlm_gepa.runtime.adapter import reflective_record - from rlm_gepa.schema import RLMGepaExampleResult - - backend = FinalBackend() - rlm = PredictRLM( - "question: str -> answer: str", - lm=MagicMock(history=[]), - execution=backend, - max_iterations=1, - verbose=False, - ) - rlm.generate_action.acall = AsyncMock( - return_value=dspy.Prediction( - reasoning="submit", - code="SUBMIT(answer=question)", - ) - ) - - prediction = rlm.forward(question="test") - result = RLMGepaExampleResult( - score=1.0, - feedback="ok", - traces=[prediction.trace], - ) - record = reflective_record(result) - evidence = record["Traces"][0]["evidence"] - kinds = [event["kind"] for event in evidence["events"]] - - assert evidence["complete"] is True - assert kinds.index("session.finalized") < kinds.index("session.released") - assert kinds.index("session.released") < kinds.index("run.succeeded") - - -def test_rlm_gepa_rejects_present_but_incomplete_strict_evidence(): - from predict_rlm.trace import RunEvidence, RunTrace - from rlm_gepa.schema import RLMGepaExampleResult, validate_example_result - - trace = RunTrace( - status="completed", - model="test", - iterations=1, - max_iterations=1, - duration_ms=1, - evidence=RunEvidence(run_id="run", complete=False), - ) - result = RLMGepaExampleResult(score=1.0, feedback="ok", traces=[trace]) - - with pytest.raises(ValueError, match="incomplete strict evidence"): - validate_example_result(result) - - -class TransformInputAdapter(InputAdapter[str]): - name = "transform" - value_type = str - - async def prepare(self, field, value, ctx): - return PreparedInput( - model_value=f"prepared:{value}", - instructions=("Use the prepared value.",), - ) - -def test_custom_contributions_drive_invocation_runtime(): +@pytest.mark.asyncio +async def test_failed_backend_exit_emits_release_failure_not_release_success(): from predict_rlm import PredictRLM - lm = MagicMock() - lm.copy.return_value = lm - lm.history = [] - backend = FinalBackend() - module_calls = 0 - - async def extra_tool(value: str) -> str: - return value - - def module(): - nonlocal module_calls - module_calls += 1 - return RuntimeContribution( - instructions=("Module instruction.",), - adapters=[TransformInputAdapter()], - tools=(CallableTool(name="extra_tool", function=extra_tool),), - packages=("module-package",), - ) - + sink = RecordingSink() + backend = FailingExitBackend() rlm = PredictRLM( "question: str -> answer: str", - lm=lm, + lm=MagicMock(history=[]), execution=backend, - modules=(module,), + events=(sink,), max_iterations=1, verbose=False, ) - rlm._build_signatures_with_files = MagicMock( - return_value=(rlm.generate_action, rlm.extract) - ) rlm.generate_action.acall = AsyncMock( return_value=dspy.Prediction( reasoning="submit", @@ -2334,164 +975,91 @@ def module(): ) ) - result = rlm.forward(question="value") + with pytest.raises(OSError, match="release failed"): + await rlm.aforward(question="test") - assert module_calls == 1 - assert result.answer == "prepared:value" - assert backend.spec.packages == ("module-package",) - assert "extra_tool" in backend.spec.tools - assert "Module instruction." in rlm.runtime_spec.instructions + kinds = [event.kind for event in sink.events] + assert RunEventKind.SESSION_RELEASED not in kinds + assert RunEventKind.SESSION_RELEASE_FAILED in kinds -def test_module_contributions_build_scalar_baseline_predictors(): +@pytest.mark.asyncio +async def test_failed_backend_exit_does_not_mask_primary_execution_error(): from predict_rlm import PredictRLM - async def module_tool(value: str) -> str: - """Return the module value unchanged.""" - return value - - backend = FinalBackend() + backend = FailingExitBackend() rlm = PredictRLM( "question: str -> answer: str", lm=MagicMock(history=[]), execution=backend, - modules=( - RuntimeContribution( - instructions=("Always use the module tool.",), - tools=( - CallableTool( - name="module_tool", - function=module_tool, - description="Return the module value unchanged.", - ), - ), - ), - ), max_iterations=1, verbose=False, ) + rlm.generate_action.acall = AsyncMock(side_effect=ValueError("primary execution failure")) - action_instructions = rlm.generate_action.signature.instructions - extract_instructions = rlm.extract.signature.instructions - - assert "Always use the module tool." in action_instructions - assert "Always use the module tool." in extract_instructions - assert "module_tool" in action_instructions - assert "module_tool" in extract_instructions - - -@pytest.mark.parametrize( - "signature", - [ - FileUnionInputSignature, - NestedFileInputSignature, - TupleFileInputSignature, - SetFileInputSignature, - DictFileInputSignature, - SequenceFileInputSignature, - NestedGenericFileInputSignature, - WorkspaceUnionInputSignature, - NestedWorkspaceInputSignature, - TupleWorkspaceInputSignature, - SetWorkspaceInputSignature, - DictWorkspaceInputSignature, - SequenceWorkspaceInputSignature, - NestedGenericWorkspaceInputSignature, - ], -) -def test_unsupported_compatibility_input_shapes_fail_before_backend_start(signature): - from predict_rlm import PredictRLM - - backend = FinalBackend() - - with pytest.raises(ValueError, match="unsupported.*annotation|Unsupported.*annotation"): - PredictRLM(signature, lm=MagicMock(history=[]), execution=backend) + with pytest.raises(ValueError, match="primary execution failure") as raised: + await rlm.aforward(question="test") - assert backend.spec is None + assert isinstance(raised.value.session_release_error, OSError) -def test_workspace_output_fails_before_backend_start(): +@pytest.mark.asyncio +async def test_cancellation_waits_for_owned_finalization_before_backend_release(): from predict_rlm import PredictRLM - class WorkspaceOutputSignature(dspy.Signature): - question: str = dspy.InputField() - workspace: Workspace = dspy.OutputField() - - backend = FinalBackend() - - with pytest.raises(ValueError, match="Workspace.*input-only"): - PredictRLM( - WorkspaceOutputSignature, - lm=MagicMock(history=[]), - execution=backend, + backend = BlockingFinalizeBackend() + rlm = PredictRLM( + "question: str -> answer: str", + lm=MagicMock(history=[]), + execution=backend, + max_iterations=1, + verbose=False, + ) + rlm.generate_action.acall = AsyncMock( + return_value=dspy.Prediction( + reasoning="submit", + code="SUBMIT(answer=question)", ) - - assert backend.spec is None - - -@pytest.mark.parametrize( - "annotation", - [ - tuple[RuntimeFile, ...], - set[RuntimeFile], - dict[str, RuntimeFile], - Sequence[RuntimeFile], - list[dict[str, tuple[RuntimeFile, ...]]], - ], -) -def test_unsupported_file_generic_output_shapes_fail_before_backend_start(annotation): - from predict_rlm import PredictRLM - - FileOutputSignature = type( - "FileOutputSignature", - (dspy.Signature,), - { - "__annotations__": {"question": str, "result": annotation}, - "question": dspy.InputField(), - "result": dspy.OutputField(), - }, ) - backend = FinalBackend() - with pytest.raises(ValueError, match="Unsupported File annotation"): - PredictRLM(FileOutputSignature, lm=MagicMock(history=[]), execution=backend) + invocation = asyncio.create_task(rlm.aforward(question="test")) + await backend.session.finalize_started.wait() + invocation.cancel() + await asyncio.sleep(0) + completed_before_finalize = invocation.done() + released_before_finalize = backend.released + backend.session.allow_finalize.set() + with pytest.raises(asyncio.CancelledError): + await invocation - assert backend.spec is None + assert not completed_before_finalize + assert not released_before_finalize + assert backend.session.finalized == 1 + assert backend.released -@pytest.mark.parametrize( - "annotation", - [ - tuple[Workspace, ...], - set[Workspace], - dict[str, Workspace], - Sequence[Workspace], - list[dict[str, tuple[Workspace, ...]]], - ], -) -def test_workspace_remains_input_only_inside_every_generic_output_shape(annotation): +@pytest.mark.asyncio +async def test_finalization_failure_preserves_primary_and_marks_evidence_incomplete(): from predict_rlm import PredictRLM - WorkspaceOutputSignature = type( - "WorkspaceOutputSignature", - (dspy.Signature,), - { - "__annotations__": {"question": str, "workspace": annotation}, - "question": dspy.InputField(), - "workspace": dspy.OutputField(), - }, + backend = FailingFinalizeBackend() + rlm = PredictRLM( + "question: str -> answer: str", + lm=MagicMock(history=[]), + execution=backend, + max_iterations=1, + verbose=False, ) + rlm.generate_action.acall = AsyncMock(side_effect=ValueError("primary failure")) - backend = FinalBackend() - - with pytest.raises(ValueError, match="Workspace.*input-only"): - PredictRLM( - WorkspaceOutputSignature, - lm=MagicMock(history=[]), - execution=backend, - ) + with pytest.raises(ValueError, match="primary failure") as raised: + await rlm.aforward(question="test") - assert backend.spec is None + assert isinstance(raised.value.session_finalize_error, OSError) + assert raised.value.trace.evidence.complete is False + assert "session.finalize_failed" in { + event.kind for event in raised.value.trace.evidence.events + } @pytest.mark.asyncio @@ -2523,9 +1091,7 @@ def acquire(spec, ctx): second.mkdir() first_ctx = RunContext(make_spec(), {}) first_mount = HostDirectoryMount(str(first), "/workspace") - async with backend.start( - ExecutionSpec(host_directory_mounts=(first_mount,)), first_ctx - ): + async with backend.start(ExecutionSpec(host_directory_mounts=(first_mount,)), first_ctx): pass second_ctx = RunContext(make_spec(), {}) @@ -2542,202 +1108,6 @@ def acquire(spec, ctx): assert acquisitions == 1 -def test_output_adapter_ambiguity_fails_before_backend_start(): - from predict_rlm import PredictRLM - - class StringOutputAdapter(OutputAdapter[str]): - value_type = str - - def __init__(self, name: str) -> None: - self.name = name - - async def reserve(self, field, value, ctx, session): - raise NotImplementedError - - async def materialize(self, reservation, submitted_value, ctx, session): - raise NotImplementedError - - lm = MagicMock() - lm.copy.return_value = lm - lm.history = [] - backend = FinalBackend() - - with pytest.raises(ValueError, match="Multiple output adapters"): - PredictRLM( - "question: str -> answer: str", - lm=lm, - execution=backend, - adapters=[StringOutputAdapter("first"), StringOutputAdapter("second")], - ) - - assert backend.spec is None - - -@pytest.mark.asyncio -async def test_kernel_cancels_and_finalizes_session_on_failure(): - from predict_rlm import PredictRLM - - lm = MagicMock() - lm.copy.return_value = lm - lm.history = [] - backend = FinalBackend() - rlm = PredictRLM( - "question -> answer", - lm=lm, - execution=backend, - max_iterations=1, - verbose=False, - ) - rlm.generate_action.acall = AsyncMock(side_effect=RuntimeError("action failed")) - rlm._build_signatures_with_files = MagicMock( - return_value=(rlm.generate_action, rlm.extract) - ) - - with pytest.raises(RuntimeError, match="action failed"): - await rlm.aforward(question="test") - - assert backend.session.cancelled == 1 - assert backend.session.finalized == 1 - - -def test_explicit_backend_mounts_and_collects_file_artifacts(tmp_path: Path): - from predict_rlm import File, PredictRLM - - source = tmp_path / "source.txt" - source.write_text("source", encoding="utf-8") - lm = MagicMock() - lm.copy.return_value = lm - lm.history = [] - backend = FinalBackend() - backend.session.final_payload = { - "result": "/sandbox/output/result/generated.txt" - } - rlm = PredictRLM( - KernelFileSignature, - lm=lm, - execution=backend, - output_dir=tmp_path / "outputs", - max_iterations=1, - verbose=False, - ) - rlm._build_signatures_with_files = MagicMock( - return_value=(rlm.generate_action, rlm.extract) - ) - rlm._prepare_file_io = MagicMock( - side_effect=AssertionError("the kernel must not build a legacy file plan") - ) - rlm.generate_action.acall = AsyncMock( - return_value=dspy.Prediction( - reasoning="write", - code="SUBMIT(result='/sandbox/output/result/generated.txt')", - ) - ) - - prediction = rlm.forward(source=File(path=str(source))) - - assert prediction.result.path.endswith("result/generated.txt") - assert Path(prediction.result.path).read_text(encoding="utf-8") == "generated" - assert len(backend.session.mounted) == 2 - assert {artifact.kind for artifact in backend.session.mounted} == { - "runtime.path", - "compat.output.directory", - } - rlm._prepare_file_io.assert_not_called() - - -def test_synced_file_operation_is_portable_to_custom_final_backend(): - from predict_rlm import PredictRLM - - backend = SyncedFinalBackend() - - def mutate(path: Annotated[Path, SyncedFile()]) -> str: - backend.session.captured_host_path = Path(path) - assert backend.session.captured_host_path.read_text(encoding="utf-8") == "before" - backend.session.captured_host_path.write_text("after", encoding="utf-8") - return "mutated" - - rlm = PredictRLM( - "question: str -> answer: str", - lm=MagicMock(history=[]), - execution=backend, - tools={"mutate": mutate}, - max_iterations=1, - verbose=False, - ) - rlm.generate_action.acall = AsyncMock( - return_value=dspy.Prediction(reasoning="mutate", code="mutate work") - ) - - result = rlm.forward(question="test") - - assert result.answer == "test" - assert backend.session.files["/sandbox/work.txt"] == "after" - assert len(backend.session.synced_mounts) == 1 - assert backend.session.captured_host_path is not None - assert not backend.session.captured_host_path.exists() - - -@pytest.mark.parametrize( - "backend_name", - [ - "jspi", - pytest.param("sbx", marks=pytest.mark.sbx), - pytest.param("sbx-pool", marks=pytest.mark.sbx), - ], -) -def test_synced_file_operation_runs_on_maintained_final_backend_lifecycles( - backend_name: str, - monkeypatch, -): - from predict_rlm import PredictRLM - from predict_rlm.backends.jspi import execution as jspi_execution - - interpreter = MaintainedSyncedInterpreter() - - def build_interpreter(**kwargs): - interpreter.tools = kwargs["tools"] - return interpreter - - kwargs = {} - pool = None - if backend_name == "jspi": - monkeypatch.setattr(jspi_execution, "JspiBackend", build_interpreter) - elif backend_name == "sbx": - from predict_rlm.backends.sbx import execution as sbx_execution - - monkeypatch.setattr(sbx_execution, "SbxBackend", build_interpreter) - kwargs["sandbox_backend"] = "sbx" - else: - pool = MaintainedSyncedPool(interpreter) - kwargs.update(sandbox_backend="sbx", sbx_pool=pool) - - def mutate(path: Annotated[Path, SyncedFile()]) -> str: - file_path = Path(path) - file_path.write_text("after", encoding="utf-8") - return "mutated" - - rlm = PredictRLM( - "question: str -> answer: str", - lm=MagicMock(history=[]), - tools={"mutate": mutate}, - max_iterations=1, - verbose=False, - **kwargs, - ) - rlm.generate_action.acall = AsyncMock( - return_value=dspy.Prediction(reasoning="mutate", code="mutate work") - ) - - result = rlm.forward(question="test") - - assert result.answer == "done" - assert interpreter.files["/sandbox/work.txt"] == "after" - if pool is None: - assert interpreter.shutdown_calls == 1 - else: - assert pool.released - - @pytest.mark.sbx def test_sync_forward_awaits_owned_sbx_host_retirement_before_loop_teardown( monkeypatch, @@ -2798,91 +1168,6 @@ async def shutdown(): assert shutdown_saw_cleanup -def test_synced_file_operation_preserves_read_only_and_custom_host_dir(tmp_path: Path): - from predict_rlm import PredictRLM - - backend = SyncedFinalBackend() - backend.session.tool_name = "inspect_file" - host_dir = tmp_path / "synced" - - def inspect_file(path: Path) -> str: - assert Path(path).parent == host_dir - Path(path).write_text("host-only", encoding="utf-8") - return "inspected" - - inspect_file.__annotations__["path"] = Annotated[ - Path, - SyncedFile(writeback=False, host_dir=str(host_dir)), - ] - - rlm = PredictRLM( - "question: str -> answer: str", - lm=MagicMock(history=[]), - execution=backend, - tools={"inspect_file": inspect_file}, - max_iterations=1, - verbose=False, - ) - rlm.generate_action.acall = AsyncMock( - return_value=dspy.Prediction(reasoning="inspect", code="inspect work") - ) - - rlm.forward(question="test") - - assert backend.session.files["/sandbox/work.txt"] == "before" - assert backend.session.synced_mounts == [] - assert (host_dir / "work.txt").read_text(encoding="utf-8") == "host-only" - assert str(host_dir) in backend.spec.extra_write_paths - - -def test_synced_file_operation_cleans_temporary_file_after_tool_failure(): - from predict_rlm import PredictRLM - - backend = SyncedFinalBackend() - captured = None - - def fail(path: Annotated[Path, SyncedFile()]) -> str: - nonlocal captured - captured = Path(path) - raise RuntimeError("tool failed") - - backend.session.tool_name = "fail" - rlm = PredictRLM( - "question: str -> answer: str", - lm=MagicMock(history=[]), - execution=backend, - tools={"fail": fail}, - max_iterations=1, - verbose=False, - ) - rlm.generate_action.acall = AsyncMock( - return_value=dspy.Prediction(reasoning="fail", code="fail work") - ) - rlm.extract.acall = AsyncMock(return_value=dspy.Prediction(answer="fallback")) - - result = rlm.forward(question="test") - - assert result.answer == "fallback" - assert captured is not None and not captured.exists() - assert backend.session.synced_mounts == [] - - -def test_sync_leaf_classification_survives_synced_file_and_evidence_wrappers(): - from predict_rlm import PredictRLM - - def inspect_file(path: Annotated[Path, SyncedFile()]) -> str: - return Path(path).name - - synced = SyncedFileToolOperation().apply( - CallableTool(name="inspect_file", function=inspect_file) - ).function - owner = type("EvidenceOwner", (), {"_evidence": lambda self: None})() - wrapped = PredictRLM._wrap_evidence_tool(owner, "inspect_file", synced) - - assert callable_has_sync_leaf(synced) - assert callable_has_sync_leaf(wrapped) - - @pytest.mark.asyncio async def test_sync_tool_cancellation_holds_custom_final_backend_lease_until_worker_stops(): from predict_rlm import PredictRLM @@ -2905,7 +1190,6 @@ def block(path: Annotated[Path, SyncedFile(writeback=False)]) -> str: class BlockingToolBackend(FinalBackend): def __init__(self) -> None: self.session = SyncedFinalSession() - self.session.tool_name = "block" self.spec = None @asynccontextmanager @@ -2995,9 +1279,7 @@ def mutate(path: Annotated[Path, SyncedFile()]) -> str: name="legacy", ownership=SessionOwnership.INJECTED, ) - wrapped = SyncedFileToolOperation().apply( - CallableTool(name="mutate", function=mutate) - ) + wrapped = SyncedFileToolOperation().apply(CallableTool(name="mutate", function=mutate)) ctx = RunContext(spec=make_spec(), input_values={}) ctx.session = session @@ -3019,40 +1301,3 @@ def mutate(path: Annotated[Path, SyncedFile()]) -> str: assert not path_removed_while_worker_live assert path_survived_worker assert captured_path is not None and not captured_path.exists() - - -def test_file_output_destination_is_not_exposed_as_sandbox_input(tmp_path: Path): - from predict_rlm import File, PredictRLM - - source = tmp_path / "source.txt" - source.write_text("source", encoding="utf-8") - backend = FinalBackend() - backend.session.final_payload = { - "result": "/sandbox/output/result/generated.txt" - } - lm = MagicMock() - lm.copy.return_value = lm - lm.history = [] - rlm = PredictRLM( - KernelFileSignature, - lm=lm, - execution=backend, - max_iterations=1, - verbose=False, - ) - rlm._build_signatures_with_files = MagicMock( - return_value=(rlm.generate_action, rlm.extract) - ) - rlm.generate_action.acall = AsyncMock( - return_value=dspy.Prediction( - reasoning="write", - code="SUBMIT(result='/sandbox/output/result/generated.txt')", - ) - ) - - rlm.forward( - source=File(path=str(source)), - result=File(path=str(tmp_path / "destination")), - ) - - assert "result" not in backend.session.variables diff --git a/tests/test_supervisor_client.py b/tests/test_supervisor_client.py index cf12b93a..cdfe25df 100644 --- a/tests/test_supervisor_client.py +++ b/tests/test_supervisor_client.py @@ -32,9 +32,7 @@ def read(self) -> str: class FakeProcess: def __init__(self, stdout_lines: list[dict[str, Any]] | None = None) -> None: self.stdin = FakePipe() - self.stdout = FakePipe( - [json.dumps(line) + "\n" for line in (stdout_lines or [])] - ) + self.stdout = FakePipe([json.dumps(line) + "\n" for line in (stdout_lines or [])]) self.stderr = FakePipe() self.returncode: int | None = None self.killed = False @@ -134,18 +132,6 @@ def _raise_execute_error(self, response: dict[str, Any]) -> None: raise CodeInterpreterError(str(error.get("message") or "runner error")) -def test_supervisor_client_discards_stale_response_then_returns_fresh() -> None: - process = FakeProcess( - [ - {"jsonrpc": "2.0", "id": 99, "result": {"output": "stale"}}, - {"jsonrpc": "2.0", "id": 1, "result": {"output": "fresh"}}, - ] - ) - client = FakeClient([process]) - - assert client.execute("print('fresh')") == "fresh" - - def test_supervisor_client_discards_stale_error_then_returns_fresh() -> None: process = FakeProcess( [ @@ -174,43 +160,3 @@ def test_supervisor_client_exhausted_stale_resync_raises_cleanly() -> None: with pytest.raises(CodeInterpreterError, match="stale.*resyncing"): client.execute("print('fresh')") - - -def test_supervisor_client_recovers_dead_supervisor_after_structured_timeout() -> None: - first = FakeProcess( - [ - { - "jsonrpc": "2.0", - "id": 1, - "result": { - "timeout": {"seconds": 0.2}, - "stdout": "before\n", - "stderr": "timed out\n", - }, - } - ] - ) - restarted = FakeProcess([]) - client = FakeClient([first, restarted]) - - timeout_result = client.execute("slow()", timeout=0.2) - first.returncode = 137 - restart_result = client.execute("next()", timeout=0.2) - - assert timeout_result.timeout_seconds == 0.2 - assert client.started == 2 - assert "fake restart diagnostic" in restart_result - assert "supervisor_returncode=137" in restart_result - assert "previous_response=structured_timeout" in restart_result - assert restarted.stdin.writes == [] - - -def test_supervisor_client_host_timeout_unwraps_recoverable_timeout() -> None: - process = FakeProcess([]) - client = FakeClient([process]) - - result = client.execute("silent()", timeout=0.2) - - assert process.killed is True - assert result.timeout_seconds == 0.2 - assert "fake supervisor restarted" in result.stderr diff --git a/tests/test_telemetry.py b/tests/test_telemetry.py index 437ed910..184b8218 100644 --- a/tests/test_telemetry.py +++ b/tests/test_telemetry.py @@ -3,42 +3,34 @@ import pytest +from predict_rlm.debug import debug_event, reset_debug_logger_for_tests from predict_rlm.telemetry import ( JsonlTelemetrySink, - NoopTelemetrySink, TelemetryContext, candidate_hash, - make_span_record, - redact_attributes, ) def test_jsonl_sink_creates_parent_directory_and_writes_one_object_per_line(tmp_path: Path): path = tmp_path / "run" / "telemetry" / "events.jsonl" sink = JsonlTelemetrySink(path) - context = TelemetryContext( - sink=sink, - trace_id="trace-1", - parent_span_id="parent-1", - run_id="run-1", - eval_kind="valset", - eval_idx=2, - attempt_id="0002", - example_id="581-46", - case_idx=56, - candidate_hash="cand_sha256_abc", - ) + context = TelemetryContext(sink=sink, trace_id="trace-1", parent_span_id="parent-1") - first = context.write_span( + context.write_span( "sandbox.execute", event_domain="sandbox", span_id="span-1", start_time_unix_nano=1_000_000_000, end_time_unix_nano=1_250_000_000, - status={"code": "ERROR", "message": "Sandbox exec timed out after 250ms"}, - attributes={"failure.class": "sandbox_exec_timeout", "process.pid": 123}, + status={"code": "ERROR", "message": "OPENAI_API_KEY=secret-value failed"}, + attributes={ + "failure.class": "sandbox_exec_timeout", + "db.password": "plain-password", + "nested": {"message": "Bearer secret-token", "safe": "ok"}, + "list": ["GITHUB_TOKEN=ghp_realvalue", "safe"], + }, ) - second = context.write_span( + context.write_span( "spreadbench.case", event_domain="spreadbench", span_id="span-2", @@ -46,80 +38,20 @@ def test_jsonl_sink_creates_parent_directory_and_writes_one_object_per_line(tmp_ status="OK", ) - assert path.parent.is_dir() - lines = path.read_text(encoding="utf-8").splitlines() - assert len(lines) == 2 - assert [json.loads(line)["span_id"] for line in lines] == ["span-1", "span-2"] - - assert first["duration_ms"] == 250 - assert first["schema_version"] == 1 - assert first["record_type"] == "span" - assert first["event_domain"] == "sandbox" - assert first["trace_id"] == "trace-1" - assert first["span_id"] == "span-1" + serialized = path.read_text(encoding="utf-8") + first, second = [json.loads(line) for line in serialized.splitlines()] + assert [first["span_id"], second["span_id"]] == ["span-1", "span-2"] + assert first["trace_id"] == second["trace_id"] == "trace-1" assert first["parent_span_id"] == "parent-1" - assert first["name"] == "sandbox.execute" - assert first["span_kind"] == "internal" - assert first["start_time_unix_nano"] == 1_000_000_000 - assert first["end_time_unix_nano"] == 1_250_000_000 - assert first["status"]["code"] == "ERROR" - assert first["attributes"]["rlm.run_id"] == "run-1" - assert first["attributes"]["rlm.eval_kind"] == "valset" - assert first["attributes"]["rlm.eval_idx"] == 2 - assert first["attributes"]["rlm.attempt_id"] == "0002" - assert first["attributes"]["spreadbench.example_id"] == "581-46" - assert first["attributes"]["spreadbench.case_idx"] == 56 - assert first["attributes"]["rlm.candidate_hash"] == "cand_sha256_abc" - assert first["attributes"]["failure.class"] == "sandbox_exec_timeout" - + assert first["duration_ms"] == 250 assert second["duration_ms"] == 5 - - -def test_make_span_record_can_override_parent_and_redacts_status_message(): - context = TelemetryContext( - sink=NoopTelemetrySink(), - trace_id="trace-1", - parent_span_id="parent-1", - ) - - record = make_span_record( - context, - name="host_tool.recalculate", - event_domain="host_tool", - span_id="span-1", - parent_span_id="override-parent", - status={"code": "ERROR", "message": "OPENAI_API_KEY=secret-value failed"}, - attributes={"safe": "value"}, - ) - - assert record["parent_span_id"] == "override-parent" - assert record["status"]["message"] == "OPENAI_API_KEY=[REDACTED] failed" - - -def test_noop_sink_and_disabled_jsonl_sink_do_not_write(tmp_path: Path): - path = tmp_path / "telemetry" / "events.jsonl" - NoopTelemetrySink().write({"trace_id": "trace-1"}) - - sink = JsonlTelemetrySink(path, enabled=False) - sink.write({"trace_id": "trace-1"}) - - assert not path.exists() - - -def test_jsonl_sink_write_errors_are_best_effort(monkeypatch: pytest.MonkeyPatch, tmp_path: Path): - path = tmp_path / "telemetry" / "events.jsonl" - sink = JsonlTelemetrySink(path) - - def fail_open(*args, **kwargs): - raise OSError("disk full") - - monkeypatch.setattr(Path, "open", fail_open) - - try: - raise RuntimeError("caller failure") - except RuntimeError as exc: - sink.write({"trace_id": "trace-1", "bad": object()}) - assert str(exc) == "caller failure" + assert first["status"]["message"] == "OPENAI_API_KEY=[REDACTED] failed" + assert first["attributes"]["failure.class"] == "sandbox_exec_timeout" + assert first["attributes"]["db.password"] == "[REDACTED]" + assert first["attributes"]["nested"] == {"message": "[REDACTED]", "safe": "ok"} + assert first["attributes"]["list"] == ["GITHUB_TOKEN=[REDACTED]", "safe"] + for secret in ("secret-value", "plain-password", "secret-token", "ghp_realvalue"): + assert secret not in serialized def test_candidate_hash_is_deterministic_for_equivalent_dicts(): @@ -127,24 +59,41 @@ def test_candidate_hash_is_deterministic_for_equivalent_dicts(): right = {"a": {"nested": True}, "b": [2, 3]} assert candidate_hash(left) == candidate_hash(right) - assert candidate_hash(left).startswith("cand_sha256_") assert candidate_hash(left) != candidate_hash({"a": {"nested": False}, "b": [2, 3]}) -def test_redact_attributes_removes_obvious_secret_values(): - redacted = redact_attributes( - { - "OPENAI_API_KEY": "sk-real-value", - "db.password": "plain-password", - "message": "export GITHUB_TOKEN=ghp_realvalue and continue", - "nested": {"auth_header": "Bearer secret-token", "safe": "ok"}, - "list": ["OPENAI_API_KEY=abc123", "safe"], - } +@pytest.fixture +def reset_debug_logging(monkeypatch): + for name in ( + "PREDICT_RLM_DEBUG", + "RLM_DEBUG", + "PREDICT_RLM_DEBUG_LOG", + "PREDICT_RLM_DEBUG_JSON", + ): + monkeypatch.delenv(name, raising=False) + reset_debug_logger_for_tests() + yield + reset_debug_logger_for_tests() + + +def test_json_debug_logging_redacts_obvious_secrets(monkeypatch, tmp_path, reset_debug_logging): + log_path = tmp_path / "predict-rlm-debug.jsonl" + monkeypatch.setenv("PREDICT_RLM_DEBUG", "1") + monkeypatch.setenv("PREDICT_RLM_DEBUG_LOG", str(log_path)) + monkeypatch.setenv("PREDICT_RLM_DEBUG_JSON", "1") + + debug_event( + "predict_rlm.redact", + api_key="sk-testsecret123456", + nested={"authorization": "Bearer abcdefghijk"}, + harmless="visible", + value="Bearer value-secret", ) - assert redacted["OPENAI_API_KEY"] == "[REDACTED]" - assert redacted["db.password"] == "[REDACTED]" - assert redacted["message"] == "export GITHUB_TOKEN=[REDACTED] and continue" - assert redacted["nested"]["auth_header"] == "[REDACTED]" - assert redacted["nested"]["safe"] == "ok" - assert redacted["list"] == ["OPENAI_API_KEY=[REDACTED]", "safe"] + record = json.loads(log_path.read_text()) + assert record["api_key"] == "[REDACTED]" + assert record["nested"]["authorization"] == "[REDACTED]" + assert record["harmless"] == "visible" + assert record["value"] == "[REDACTED]" + assert "sk-testsecret123456" not in log_path.read_text() + assert "Bearer abcdefghijk" not in log_path.read_text() diff --git a/tests/test_telemetry_analyzer.py b/tests/test_telemetry_analyzer.py index a17c260c..536f94c5 100644 --- a/tests/test_telemetry_analyzer.py +++ b/tests/test_telemetry_analyzer.py @@ -1,48 +1,11 @@ import pytest -from predict_rlm.telemetry import classify_failure, classify_zero_score_failure +from predict_rlm.telemetry import classify_failure from rlm_gepa.runtime.telemetry_analyzer import analyze_run, analyze_trace_rows pytestmark = pytest.mark.gepa -@pytest.mark.parametrize( - ("failure_class", "expected"), - [ - ("sandbox_lifecycle_failure", "sandbox_lifecycle_failure"), - ("rlm_iteration_execution_timeout", "rlm_iteration_execution_timeout"), - ("sandbox_exec_timeout", "sandbox_exec_timeout"), - ("host_tool_timeout_or_leak", "host_tool_timeout_or_leak"), - ("outer_task_timeout", "outer_task_timeout"), - ("evaluator_limitation", "evaluator_limitation"), - ("evaluator_exception", "evaluator_exception"), - ("model_output_truncated", "model_output_truncated"), - ("model_no_code_generated", "model_no_code_generated"), - ("model_generated_bad_code", "model_generated_bad_code"), - ("resource_saturation_unknown", "resource_saturation_unknown"), - ], -) -def test_classifies_closed_failure_classes_from_failure_class_attribute( - failure_class: str, - expected: str, -): - events = [ - { - "name": "synthetic.event", - "status": {"code": "ERROR"}, - "attributes": {"failure.class": failure_class}, - } - ] - - assert classify_failure({"score": 0}, events) == expected - - -def test_missing_or_partial_evidence_returns_unknown(): - assert classify_failure({"score": 0}, []) == "unknown" - assert classify_failure({"score": 0, "status": "failed"}, None) == "unknown" - assert classify_failure(None, [{"name": "sandbox.execute", "status": {"code": "OK"}}]) == "unknown" - - def test_precedence_prefers_sandbox_lifecycle_over_all_other_classes(): events = [ {"attributes": {"failure.class": "model_generated_bad_code"}}, @@ -58,64 +21,6 @@ def test_precedence_prefers_sandbox_lifecycle_over_all_other_classes(): assert classify_failure({"score": 0}, events) == "sandbox_lifecycle_failure" -def test_precedence_prefers_rlm_iteration_timeout_over_sandbox_exec_timeout(): - events = [ - {"attributes": {"failure.class": "sandbox_exec_timeout"}}, - {"attributes": {"failure.class": "rlm_iteration_execution_timeout"}}, - ] - - assert classify_failure({"score": 0}, events) == "rlm_iteration_execution_timeout" - - -def test_precedence_prefers_sandbox_exec_timeout_over_host_tool_and_outer_timeout(): - events = [ - {"attributes": {"failure.class": "outer_task_timeout"}}, - {"attributes": {"failure.class": "host_tool_timeout_or_leak"}}, - {"attributes": {"failure.class": "sandbox_exec_timeout"}}, - ] - - assert classify_failure({"score": 0}, events) == "sandbox_exec_timeout" - - -def test_precedence_prefers_host_tool_timeout_over_outer_timeout(): - events = [ - {"attributes": {"failure.class": "outer_task_timeout"}}, - {"attributes": {"failure.class": "host_tool_timeout_or_leak"}}, - ] - - assert classify_failure({"score": 0}, events) == "host_tool_timeout_or_leak" - - -def test_outer_task_timeout_is_used_when_no_lower_level_timeout_evidence_exists(): - events = [ - { - "name": "gepa.case.outer_timeout", - "status": {"code": "ERROR", "message": "outer task timed out after 300s"}, - "attributes": {"failure.class": "outer_task_timeout"}, - } - ] - - assert classify_failure({"score": 0}, events) == "outer_task_timeout" - - -def test_evaluator_limitation_precedes_evaluator_exception(): - events = [ - {"attributes": {"failure.class": "evaluator_exception"}}, - {"attributes": {"failure.class": "evaluator_limitation"}}, - ] - - assert classify_failure({"score": 0}, events) == "evaluator_limitation" - - -def test_model_no_code_precedes_model_generated_bad_code(): - events = [ - {"attributes": {"failure.class": "model_generated_bad_code"}}, - {"attributes": {"failure.class": "model_no_code_generated"}}, - ] - - assert classify_failure({"score": 0}, events) == "model_no_code_generated" - - def test_model_output_truncated_precedes_generic_no_code_failure(): events = [ { @@ -135,12 +40,6 @@ def test_model_output_truncated_precedes_generic_no_code_failure(): assert classify_failure({"score": 0}, events) == "model_output_truncated" -def test_resource_saturation_is_only_above_unknown(): - events = [{"attributes": {"failure.class": "resource_saturation_unknown"}}] - - assert classify_failure({"score": 0}, events) == "resource_saturation_unknown" - - def test_row_failure_class_participates_in_precedence_with_events(): row = {"score": 0, "failure_class": "outer_task_timeout"} events = [{"attributes": {"failure.class": "sandbox_exec_timeout"}}] @@ -148,25 +47,6 @@ def test_row_failure_class_participates_in_precedence_with_events(): assert classify_failure(row, events) == "sandbox_exec_timeout" -def test_infers_classes_from_otel_shaped_synthetic_events(): - events = [ - { - "name": "sandbox.health_check", - "event_domain": "sandbox", - "status": {"code": "ERROR", "message": "No response during health check"}, - "attributes": {}, - }, - { - "name": "host_tool.recalculate", - "event_domain": "host_tool", - "status": {"code": "ERROR", "message": "LibreOffice subprocess timed out"}, - "attributes": {}, - }, - ] - - assert classify_zero_score_failure({"score": 0}, events) == "sandbox_lifecycle_failure" - - def test_analyzer_loads_task_traces_and_telemetry_events(tmp_path): trace_dir = tmp_path / "task_traces" telemetry_dir = tmp_path / "telemetry" diff --git a/tests/test_terminal_bench_web_search.py b/tests/test_terminal_bench_web_search.py deleted file mode 100644 index 14c27b05..00000000 --- a/tests/test_terminal_bench_web_search.py +++ /dev/null @@ -1,109 +0,0 @@ -from __future__ import annotations - -import asyncio -import json -import sys -from pathlib import Path - -ROOT = Path(__file__).resolve().parents[1] -TERMINAL_BENCH_DIR = ROOT / "examples" / "terminal_bench" -if str(TERMINAL_BENCH_DIR) not in sys.path: - sys.path.insert(0, str(TERMINAL_BENCH_DIR)) - -from terminal_bench_rlm.skills import ( # noqa: E402 - DEFAULT_TERMINAL_BENCH_SKILL_INSTRUCTIONS, - build_terminal_bench_skill, -) -from terminal_bench_rlm.tools.tbench_agent import _REMOTE_CONTROLLER_ENV_KEYS # noqa: E402 -from terminal_bench_rlm.web_search import web_search # noqa: E402 - - -class DummySkill: - def __init__(self, *, name, instructions, tools): - self.name = name - self.instructions = instructions - self.tools = tools - - -class FakeResponse: - def __enter__(self): - return self - - def __exit__(self, *args): - return False - - def read(self): - return json.dumps( - { - "choices": [{"message": {"content": "Answer text"}}], - "citations": ["https://example.com/source"], - "search_results": [ - { - "title": "Source", - "url": "https://example.com/source", - "date": "2026-01-01", - "snippet": "Useful snippet", - "ignored": "field", - }, - {"title": "Second", "url": "https://example.com/second"}, - ], - } - ).encode() - - -def test_terminal_bench_skill_registers_web_search_tool(): - skill = build_terminal_bench_skill(DummySkill) - - assert skill.tools == {"web_search": web_search} - assert "await web_search(query, max_results=5)" in skill.instructions - assert 'domains=["example.com"]' in skill.instructions - assert "await web_search(query, max_results=5)" in DEFAULT_TERMINAL_BENCH_SKILL_INSTRUCTIONS - assert 'domains=["example.com"]' in DEFAULT_TERMINAL_BENCH_SKILL_INSTRUCTIONS - - -def test_perplexity_key_is_forwarded_to_daytona_remote_controller(): - assert "PERPLEXITY_API_KEY" in _REMOTE_CONTROLLER_ENV_KEYS - - -def test_web_search_returns_cited_json(monkeypatch): - captured = {} - - def fake_urlopen(request, timeout): - captured["timeout"] = timeout - captured["body"] = json.loads(request.data.decode()) - captured["authorization"] = request.headers["Authorization"] - return FakeResponse() - - monkeypatch.setenv("PERPLEXITY_API_KEY", "test-key") - monkeypatch.setattr("urllib.request.urlopen", fake_urlopen) - - payload = json.loads( - asyncio.run( - web_search( - "what is predict-rlm?", - max_results=1, - domains=[" example.com "], - timeout=3, - ) - ) - ) - - assert captured["timeout"] == 3 - assert captured["authorization"] == "Bearer test-key" - assert captured["body"]["model"] == "sonar" - assert captured["body"]["messages"][-1]["content"] == "what is predict-rlm?" - assert captured["body"]["search_domain_filter"] == ["example.com"] - assert payload == { - "query": "what is predict-rlm?", - "domains": ["example.com"], - "answer": "Answer text", - "citations": ["https://example.com/source"], - "search_results": [ - { - "title": "Source", - "url": "https://example.com/source", - "date": "2026-01-01", - "snippet": "Useful snippet", - } - ], - } diff --git a/tests/test_tool_call_timeout.py b/tests/test_tool_call_timeout.py index 924dc897..a4f567e0 100644 --- a/tests/test_tool_call_timeout.py +++ b/tests/test_tool_call_timeout.py @@ -1,35 +1,9 @@ -"""RED-GREEN repro for unbounded host-side tool calls. - -Background: - Sandbox code calling ``await recalculate(path)`` triggers a host - tool dispatch in ``JspiBackend._execute_tool_async``. If the - tool is slow (e.g. LibreOffice on a whole-column formula → 2-3 - minutes), the overall ``execute`` round-trip blows past the - ``_exec_timeout`` ceiling. That timeout **kills the Deno - subprocess**, raises ``SandboxFatalError``, and turns what should - have been a recoverable tool error into a cascade of - ``[Errno 9] Bad file descriptor`` retries. A 2026-04-18 gemini - eval lost 19 cases to this failure mode — every one of them - scored 0 by the time ``task_timeout=600s`` finally fired. - - The fix: give each tool call its own wall-clock budget - (``TOOL_CALL_TIMEOUT_SEC``, default 180s) via ``asyncio.wait_for``. - If the tool exceeds its budget, return a clean error response to - the sandbox — deno's ``await tool()`` resumes with the error, - exec continues, the RLM can see "[Error] tool timed out" and - rewrite its code using a different approach. The sandbox stays - alive, the case stays recoverable. - -RED: a mock tool that sleeps forever hangs ``_execute_tool_async``. -GREEN: it returns an error response with a timeout message within - ~TOOL_CALL_TIMEOUT_SEC. -""" +"""Host tool deadlines remain recoverable for async, sync, and wrapped tools.""" from __future__ import annotations import asyncio import concurrent.futures -import inspect import time import pytest @@ -87,50 +61,13 @@ async def _run(): f"tool returned but took {elapsed:.2f}s — expected ~0.3s based on " f"the monkeypatched timeout" ) - assert "error" in response, ( - f"expected error response after timeout, got {response!r}" - ) + assert "error" in response, f"expected error response after timeout, got {response!r}" err = str(response.get("error") or "") assert "timed out" in err.lower() or "timeout" in err.lower(), ( f"error message should mention the timeout; got {err!r}" ) -def test_async_tool_that_completes_quickly_is_not_affected(monkeypatch): - """Guardrail: normal fast tools must continue to return their - results unchanged — the timeout is a ceiling, not a delay. - """ - monkeypatch.setattr(rlm_interpreter, "TOOL_CALL_TIMEOUT_SEC", 1.0) - - async def _fast_tool(**_kwargs): - return "ok" - - interp = _build_interp_with_tool(_fast_tool) - response = asyncio.run( - interp._execute_tool_async("slow_tool", {"args": [], "kwargs": {}}) - ) - assert response.get("value") == "ok" - assert "error" not in response - - -def test_tool_exception_still_routes_through_error_path(monkeypatch): - """If a tool raises (e.g. ValueError inside the tool), the existing - ``except Exception`` in _execute_tool_async captures it and returns - ``{"error": ...}``. The timeout wrap must not change this behaviour. - """ - monkeypatch.setattr(rlm_interpreter, "TOOL_CALL_TIMEOUT_SEC", 1.0) - - async def _raising_tool(**_kwargs): - raise ValueError("tool blew up") - - interp = _build_interp_with_tool(_raising_tool) - response = asyncio.run( - interp._execute_tool_async("slow_tool", {"args": [], "kwargs": {}}) - ) - assert "error" in response - assert "blew up" in str(response["error"]) - - def test_sync_tool_timeout_does_not_poison_executor(monkeypatch): monkeypatch.setattr(rlm_interpreter, "TOOL_CALL_TIMEOUT_SEC", 0.05) @@ -144,9 +81,7 @@ def _fast_tool(): interp = _build_interp_with_tool(_slow_tool) interp._executor = concurrent.futures.ThreadPoolExecutor(max_workers=1) try: - first = asyncio.run( - interp._execute_tool_async("slow_tool", {"args": [], "kwargs": {}}) - ) + first = asyncio.run(interp._execute_tool_async("slow_tool", {"args": [], "kwargs": {}})) assert "error" in first interp.tools["slow_tool"] = _fast_tool @@ -175,8 +110,6 @@ def _blocked_tool(): wrapped = PredictRLM._wrap_evidence_tool(owner, "slow_tool", _blocked_tool) interp = _build_interp_with_tool(wrapped) - assert not inspect.iscoroutinefunction(inspect.unwrap(wrapped)) - async def _run(): started_at = time.monotonic() response = await interp._execute_tool_async( diff --git a/tests/test_trace.py b/tests/test_trace.py index 77a540c0..a016efbc 100644 --- a/tests/test_trace.py +++ b/tests/test_trace.py @@ -1,9 +1,7 @@ """Tests for structured trace output.""" -import contextvars import json -import time -from unittest.mock import MagicMock +from types import SimpleNamespace import pytest @@ -13,230 +11,27 @@ LMUsage, PredictCallDetail, PredictCallGroup, - ProposerRunTrace, RunEvidence, RunEvidenceEvent, RunTrace, TokenUsage, ToolCall, _RawPredictCall, - _sanitize_for_trace, drain_predict_calls, drain_tool_calls, init_predict_call_collector, init_tool_call_collector, lm_completion_metadata_since, lm_finish_since, - ms_since, record_predict_call, record_tool_call, reset_predict_call_collector, reset_tool_call_collector, - snapshot_lm_history_len, usage_since, ) -class TestTokenUsage: - def test_defaults_to_zero(self): - usage = TokenUsage() - assert usage.input_tokens == 0 - assert usage.output_tokens == 0 - assert usage.cost == 0.0 - - def test_iadd(self): - a = TokenUsage(input_tokens=10, output_tokens=5, cost=0.01) - b = TokenUsage(input_tokens=20, output_tokens=10, cost=0.02) - a += b - assert a.input_tokens == 30 - assert a.output_tokens == 15 - assert a.cost == pytest.approx(0.03) - - -class TestPredictCallGroup: - def test_fields(self): - group = PredictCallGroup( - signature="q -> a", - model="openai/gpt-4o", - calls=[ - PredictCallDetail(duration_ms=90, usage=TokenUsage(input_tokens=40, output_tokens=20, cost=0.004)), - PredictCallDetail(duration_ms=110, usage=TokenUsage(input_tokens=60, output_tokens=30, cost=0.006)), - ], - ) - assert group.signature == "q -> a" - assert group.model == "openai/gpt-4o" - assert len(group.calls) == 2 - assert group.calls[0].duration_ms == 90 - - -class TestIterationStep: - def test_fields(self): - step = IterationStep( - iteration=1, - reasoning="think", - code="print(1)", - output="1", - untruncated_output="1", - duration_ms=500, - ) - assert step.iteration == 1 - assert step.predict_calls == [] - - def test_output_vs_untruncated(self): - long_output = "x" * 10000 - truncated = long_output[:5000] + f"\n... (truncated to 5000/{len(long_output):,} chars)" - step = IterationStep( - iteration=1, - reasoning="", - code="print('x' * 10000)", - output=truncated, - untruncated_output=long_output, - duration_ms=100, - ) - assert len(step.untruncated_output) == 10000 - assert len(step.output) < 6000 - assert "truncated" in step.output - - class TestRunTrace: - def test_serialization(self): - trace = RunTrace( - status="completed", - model="openai/gpt-5", - sub_model="openai/gpt-4o", - iterations=2, - max_iterations=5, - duration_ms=1000, - usage=LMUsage( - main=TokenUsage(input_tokens=70, output_tokens=40, cost=0.04), - sub=TokenUsage(input_tokens=30, output_tokens=10, cost=0.01), - ), - steps=[ - IterationStep( - iteration=1, - reasoning="step 1", - code="x = 1", - output="", - untruncated_output="", - duration_ms=400, - lm=LMFinishMetadata(finish_reason="length"), - predict_calls=[ - PredictCallGroup( - signature="q -> a", - model="openai/gpt-4o", - calls=[ - PredictCallDetail( - duration_ms=200, - usage=TokenUsage(input_tokens=30, output_tokens=10, cost=0.01), - lm=LMFinishMetadata(finish_reason="length"), - ) - ], - ) - ], - ), - IterationStep( - iteration=2, - reasoning="step 2", - code="SUBMIT(x)", - output="FINAL: {'answer': 1}", - untruncated_output="FINAL: {'answer': 1}", - duration_ms=600, - ), - ], - ) - data = trace.model_dump() - assert data["status"] == "completed" - assert data["model"] == "openai/gpt-5" - assert data["sub_model"] == "openai/gpt-4o" - assert data["iterations"] == 2 - assert len(data["steps"]) == 2 - assert len(data["steps"][0]["predict_calls"]) == 1 - assert data["steps"][0]["predict_calls"][0]["signature"] == "q -> a" - assert data["steps"][0]["predict_calls"][0]["model"] == "openai/gpt-4o" - assert data["usage"]["main"]["input_tokens"] == 70 - assert data["usage"]["sub"]["input_tokens"] == 30 - - def test_exportable_json_keeps_lm_finish_metadata_compact(self): - trace = RunTrace( - status="completed", - model="openai/gpt-5", - iterations=1, - max_iterations=5, - duration_ms=100, - usage=LMUsage( - main=TokenUsage(input_tokens=100, output_tokens=50, cost=0.01), - sub=TokenUsage(input_tokens=20, output_tokens=10, cost=0.002), - ), - steps=[ - IterationStep( - iteration=1, - reasoning="think", - code="answer = predict(q='x')", - output="ok", - untruncated_output="ok", - duration_ms=100, - lm=LMFinishMetadata(finish_reason="length"), - predict_calls=[ - PredictCallGroup( - signature="q -> a", - model="openai/gpt-4o", - total_usage=TokenUsage(input_tokens=20, output_tokens=10, cost=0.002), - calls=[ - PredictCallDetail( - duration_ms=20, - usage=TokenUsage(input_tokens=20, output_tokens=10, cost=0.002), - lm=LMFinishMetadata(finish_reason="length"), - ) - ], - ) - ], - ) - ], - ) - - import json - - data = json.loads(trace.to_exportable_json()) - serialized = trace.to_exportable_json() - assert '"truncation"' not in serialized - assert data["usage"]["main"] == { - "input_tokens": 100, - "output_tokens": 50, - "cost": 0.01, - "cache_hits": 0, - } - assert data["usage"]["sub"] == { - "input_tokens": 20, - "output_tokens": 10, - "cost": 0.002, - "cache_hits": 0, - } - assert data["steps"][0]["lm"] == {"finish_reason": "length"} - call = data["steps"][0]["predict_calls"][0]["calls"][0] - assert call["lm"] == {"finish_reason": "length"} - forbidden = {"truncated", "truncation_reason", "max_tokens", "output_tokens"} - assert forbidden.isdisjoint(data["steps"][0]["lm"]) - assert forbidden.isdisjoint(call["lm"]) - - def test_sub_model_optional(self): - trace = RunTrace( - status="completed", - model="openai/gpt-5", - iterations=1, - max_iterations=5, - duration_ms=100, - ) - assert trace.sub_model is None - assert trace.usage.sub.input_tokens == 0 - - def test_status_literals(self): - for status in ("in_progress", "completed", "max_iterations", "error"): - trace = RunTrace( - status=status, model="openai/gpt-5", - iterations=1, max_iterations=5, duration_ms=100, - ) - assert trace.status == status - def test_atomic_export_replaces_in_progress_trace(self, tmp_path): from predict_rlm.predict_rlm import PredictRLM @@ -267,35 +62,33 @@ def test_atomic_export_replaces_in_progress_trace(self, tmp_path): assert json.loads(path.read_text())["status"] == "completed" assert not list(tmp_path.glob("*.tmp")) - def test_to_exportable_json_returns_string(self): - trace = RunTrace( - status="completed", model="openai/gpt-5", - iterations=1, max_iterations=5, duration_ms=100, - ) - result = trace.to_exportable_json() - assert isinstance(result, str) - import json - data = json.loads(result) - assert data["status"] == "completed" - def test_to_exportable_json_sanitizes_base64(self): b64 = "A" * 40000 trace = RunTrace( - status="completed", model="openai/gpt-5", - iterations=1, max_iterations=5, duration_ms=100, + status="completed", + model="openai/gpt-5", + iterations=1, + max_iterations=5, + duration_ms=100, steps=[ IterationStep( - iteration=1, reasoning="", code="", - output="", untruncated_output="", duration_ms=100, + iteration=1, + reasoning="", + code="", + output="", + untruncated_output="", + duration_ms=100, predict_calls=[ PredictCallGroup( signature="page: dspy.Image -> answer", model="openai/gpt-4o", - calls=[PredictCallDetail( - duration_ms=50, - input={"page": f"data:image/png;base64,{b64}"}, - output={"answer": "hello"}, - )], + calls=[ + PredictCallDetail( + duration_ms=50, + input={"page": f"data:image/png;base64,{b64}"}, + output={"answer": "hello"}, + ) + ], ) ], ) @@ -377,12 +170,6 @@ def test_to_proposer_keeps_behavioral_evidence_without_accounting(self): step = data["steps"][0] predict_call = step["predict_calls"][0]["calls"][0] - assert isinstance(proposer, ProposerRunTrace) - assert data["status"] == "completed" - assert data["model"] == "openai/gpt-5" - assert data["sub_model"] == "openai/gpt-4o" - assert data["iterations"] == 1 - assert data["max_iterations"] == 5 assert step["reasoning"] == "inspect" assert step["code"] == "answer = tool('x')" assert step["output"] == "truncated" @@ -406,53 +193,6 @@ def test_to_proposer_keeps_behavioral_evidence_without_accounting(self): for field in forbidden: assert f'"{field}"' not in serialized - def test_to_proposer_sanitizes_base64_payloads(self): - b64 = "A" * 40000 - trace = RunTrace( - status="completed", - model="openai/gpt-5", - iterations=1, - max_iterations=5, - duration_ms=100, - steps=[ - IterationStep( - iteration=1, - reasoning=f"see data:image/png;base64,{b64}", - code="SUBMIT(answer)", - output="ok", - untruncated_output="ok", - duration_ms=100, - tool_calls=[ - ToolCall( - name="tool", - args=[f"data:image/png;base64,{b64}"], - result={"image": f"data:image/png;base64,{b64}"}, - duration_ms=1, - ) - ], - predict_calls=[ - PredictCallGroup( - signature="page: dspy.Image -> answer", - model="openai/gpt-4o", - calls=[ - PredictCallDetail( - duration_ms=50, - input={"page": f"data:image/png;base64,{b64}"}, - output={"answer": "hello"}, - ) - ], - ) - ], - ) - ], - ) - - result = trace.to_proposer_json() - assert "AAAA" not in result - assert result.count("") == 4 - full = trace.model_dump() - assert b64 in full["steps"][0]["reasoning"] - def test_to_proposer_projects_strict_evidence_without_raw_or_accounting_data(self): b64 = "A" * 40000 trace = RunTrace( @@ -537,362 +277,154 @@ def test_to_proposer_projects_strict_evidence_without_raw_or_accounting_data(sel ): assert f'"{field}"' not in serialized - def test_to_exportable_json_writes_file(self, tmp_path): - trace = RunTrace( - status="completed", model="openai/gpt-5", - iterations=1, max_iterations=5, duration_ms=100, - ) - out = tmp_path / "trace.json" - result = trace.to_exportable_json(out) - assert out.exists() - assert result == out.read_text() - class TestPredictCallCollector: - def test_different_signatures_not_aggregated(self): - init_predict_call_collector() - record_predict_call(_RawPredictCall( - signature="a -> b", instructions=None, model="openai/gpt-4o", - duration_ms=50, usage=TokenUsage(), input={"a": "1"}, output={"b": "x"}, - )) - record_predict_call(_RawPredictCall( - signature="c -> d", instructions=None, model="openai/gpt-4o", - duration_ms=30, usage=TokenUsage(), input={"c": "2"}, output={"d": "y"}, - )) - groups = drain_predict_calls() - assert len(groups) == 2 - assert groups[0].signature == "a -> b" - assert groups[1].signature == "c -> d" - assert len(groups[0].calls) == 1 - assert groups[0].calls[0].output == {"b": "x"} - # Drain clears the list - assert drain_predict_calls() == [] - - def test_same_signature_aggregated(self): - init_predict_call_collector() - record_predict_call(_RawPredictCall( - signature="page: dspy.Image -> items: list[str]", - instructions="Extract items", - model="openai/gpt-4o", - duration_ms=100, - usage=TokenUsage(input_tokens=50, output_tokens=20, cost=0.01), - input={"page": "url1"}, output={"items": ["a"]}, - )) - record_predict_call(_RawPredictCall( - signature="page: dspy.Image -> items: list[str]", - instructions="Extract items", - model="openai/gpt-4o", - duration_ms=150, - usage=TokenUsage(input_tokens=60, output_tokens=30, cost=0.02), - input={"page": "url2"}, output={"items": ["b", "c"]}, - )) - record_predict_call(_RawPredictCall( - signature="page: dspy.Image -> items: list[str]", - instructions="Extract items", - model="openai/gpt-4o", - duration_ms=120, - usage=TokenUsage(input_tokens=55, output_tokens=25, cost=0.015), - input={"page": "url3"}, output={"items": []}, - )) - groups = drain_predict_calls() - assert len(groups) == 1 - agg = groups[0] - assert len(agg.calls) == 3 - assert agg.calls[0].duration_ms == 100 - assert agg.calls[0].output == {"items": ["a"]} - assert agg.calls[1].usage.input_tokens == 60 - assert agg.calls[2].output == {"items": []} - - def test_different_instructions_not_aggregated(self): - init_predict_call_collector() - record_predict_call(_RawPredictCall( - signature="q -> a", instructions="Task A", model="m", - duration_ms=10, usage=TokenUsage(), input={"q": "x"}, output={"a": "1"}, - )) - record_predict_call(_RawPredictCall( - signature="q -> a", instructions="Task B", model="m", - duration_ms=10, usage=TokenUsage(), input={"q": "y"}, output={"a": "2"}, - )) - groups = drain_predict_calls() - assert len(groups) == 2 + def test_groups_by_signature_instructions_and_model(self): + token = init_predict_call_collector() + try: + for signature, instructions, model, answer in [ + ("q -> a", "extract", "m", "first"), + ("q -> a", "extract", "m", "second"), + ("q -> a", "summarize", "m", "different instructions"), + ("q -> a", "extract", "other", "different model"), + ("text -> label", "extract", "m", "different signature"), + ]: + record_predict_call( + _RawPredictCall( + signature=signature, + instructions=instructions, + model=model, + duration_ms=10, + usage=TokenUsage(input_tokens=5, cost=0.01), + input={"q": "input"}, + output={"a": answer}, + ) + ) + groups = drain_predict_calls() + assert [[call.output["a"] for call in group.calls] for group in groups] == [ + ["first", "second"], + ["different instructions"], + ["different model"], + ["different signature"], + ] + assert groups[0].total_usage.input_tokens == 10 + assert groups[0].total_usage.cost == pytest.approx(0.02) + assert drain_predict_calls() == [] + finally: + reset_predict_call_collector(token) def test_nested_collector_restores_parent_calls(self): outer_token = init_predict_call_collector() - record_predict_call(_RawPredictCall( - signature="outer-before", instructions=None, model="m", - duration_ms=10, usage=TokenUsage(), input={}, output={}, - )) + record_predict_call( + _RawPredictCall( + signature="outer-before", + instructions=None, + model="m", + duration_ms=10, + usage=TokenUsage(), + input={}, + output={}, + ) + ) inner_token = init_predict_call_collector() - record_predict_call(_RawPredictCall( - signature="inner", instructions=None, model="m", - duration_ms=10, usage=TokenUsage(), input={}, output={}, - )) + record_predict_call( + _RawPredictCall( + signature="inner", + instructions=None, + model="m", + duration_ms=10, + usage=TokenUsage(), + input={}, + output={}, + ) + ) inner_groups = drain_predict_calls() reset_predict_call_collector(inner_token) - record_predict_call(_RawPredictCall( - signature="outer-after", instructions=None, model="m", - duration_ms=10, usage=TokenUsage(), input={}, output={}, - )) + record_predict_call( + _RawPredictCall( + signature="outer-after", + instructions=None, + model="m", + duration_ms=10, + usage=TokenUsage(), + input={}, + output={}, + ) + ) outer_groups = drain_predict_calls() reset_predict_call_collector(outer_token) assert [group.signature for group in inner_groups] == ["inner"] assert [group.signature for group in outer_groups] == ["outer-before", "outer-after"] - def test_record_without_init_is_silent(self): - from predict_rlm import trace - - original = trace._predict_calls - trace._predict_calls = contextvars.ContextVar("_predict_calls_fresh") - try: - record_predict_call(_RawPredictCall( - signature="orphan", instructions=None, model="m", - duration_ms=1, usage=TokenUsage(), input={}, output={}, - )) - assert drain_predict_calls() == [] - finally: - trace._predict_calls = original - - -class TestToolCall: - def test_fields(self): - call = ToolCall( - name="read_pdf", args=[], kwargs={"path": "/tmp/doc.pdf"}, - result='{"pages": 5}', duration_ms=200, - ) - assert call.name == "read_pdf" - assert call.error is None - assert call.result == '{"pages": 5}' - - def test_error_field(self): - call = ToolCall( - name="bad_tool", args=[], kwargs={}, - result="", error="FileNotFoundError: no such file", duration_ms=10, - ) - assert call.error == "FileNotFoundError: no such file" - - def test_serialization(self): - call = ToolCall( - name="search", args=["query"], kwargs={"limit": 10}, - result='["a", "b"]', duration_ms=50, - ) - data = call.model_dump() - assert data["name"] == "search" - assert data["args"] == ["query"] - assert data["kwargs"] == {"limit": 10} - class TestToolCallCollector: - def test_init_drain_cycle(self): - init_tool_call_collector() - record_tool_call(ToolCall( - name="tool_a", args=[], kwargs={"x": 1}, - result="ok", duration_ms=50, - )) - record_tool_call(ToolCall( - name="tool_b", args=[1, 2], kwargs={}, - result="done", duration_ms=30, - )) - calls = drain_tool_calls() - assert len(calls) == 2 - assert calls[0].name == "tool_a" - assert calls[1].name == "tool_b" - assert calls[1].args == [1, 2] - # Drain clears the list - assert drain_tool_calls() == [] - - def test_error_calls_recorded(self): - init_tool_call_collector() - record_tool_call(ToolCall( - name="failing_tool", args=[], kwargs={}, - result="", error="boom", duration_ms=5, - )) - calls = drain_tool_calls() - assert len(calls) == 1 - assert calls[0].error == "boom" - def test_nested_collector_restores_parent_calls(self): outer_token = init_tool_call_collector() - record_tool_call(ToolCall( - name="outer_before", args=[], kwargs={}, result="ok", duration_ms=1, - )) + record_tool_call( + ToolCall( + name="outer_before", + args=[], + kwargs={}, + result="ok", + duration_ms=1, + ) + ) inner_token = init_tool_call_collector() - record_tool_call(ToolCall( - name="inner", args=[], kwargs={}, result="ok", duration_ms=1, - )) + record_tool_call( + ToolCall( + name="inner", + args=[], + kwargs={}, + result="", + error="tool failed", + duration_ms=1, + ) + ) inner_calls = drain_tool_calls() + assert drain_tool_calls() == [] reset_tool_call_collector(inner_token) - record_tool_call(ToolCall( - name="outer_after", args=[], kwargs={}, result="ok", duration_ms=1, - )) + record_tool_call( + ToolCall( + name="outer_after", + args=[], + kwargs={}, + result="ok", + duration_ms=1, + ) + ) outer_calls = drain_tool_calls() reset_tool_call_collector(outer_token) assert [call.name for call in inner_calls] == ["inner"] + assert inner_calls[0].error == "tool failed" assert [call.name for call in outer_calls] == ["outer_before", "outer_after"] - def test_record_without_init_is_silent(self): - from predict_rlm import trace - - original = trace._tool_calls - trace._tool_calls = contextvars.ContextVar("_tool_calls_fresh") - try: - record_tool_call(ToolCall( - name="orphan", args=[], kwargs={}, result="", duration_ms=1, - )) - assert drain_tool_calls() == [] - finally: - trace._tool_calls = original - - -class TestSnapshotLmHistoryLen: - def test_with_history(self): - lm = MagicMock() - lm.history = [{"usage": {}}, {"usage": {}}] - assert snapshot_lm_history_len(lm) == 2 - - def test_without_history(self): - lm = MagicMock(spec=[]) - assert snapshot_lm_history_len(lm) == 0 - - def test_empty_history(self): - lm = MagicMock() - lm.history = [] - assert snapshot_lm_history_len(lm) == 0 - class TestUsageSince: - def test_sums_new_entries(self): - lm = MagicMock() - lm.history = [ - {"usage": {"prompt_tokens": 100, "completion_tokens": 50}, "cost": 0.01}, - {"usage": {"prompt_tokens": 200, "completion_tokens": 100}, "cost": 0.02}, - {"usage": {"prompt_tokens": 300, "completion_tokens": 150}, "cost": 0.03}, - ] + def test_history_delta_excludes_cached_cost_without_losing_real_usage(self): + lm = SimpleNamespace( + history=[ + {"usage": {"prompt_tokens": 900}, "cost": 1.0}, + {"usage": {"prompt_tokens": 100, "completion_tokens": 20}, "cost": 0.01}, + {"usage": {}, "cost": 0.01, "response": SimpleNamespace(cache_hit=True)}, + {"usage": {}, "cost": 0.01}, + {"usage": {"prompt_tokens": 0, "completion_tokens": 0}, "cost": 0}, + {"usage": {"prompt_tokens": 200, "completion_tokens": 30}, "cost": 0.02}, + ] + ) usage = usage_since(lm, 1) - assert usage.input_tokens == 500 - assert usage.output_tokens == 250 - assert usage.cost == pytest.approx(0.05) - - def test_since_zero_sums_all(self): - lm = MagicMock() - lm.history = [ - {"usage": {"prompt_tokens": 100, "completion_tokens": 50}, "cost": 0.01}, - ] - usage = usage_since(lm, 0) - assert usage.input_tokens == 100 + assert usage.input_tokens == 300 assert usage.output_tokens == 50 - - def test_since_beyond_length_returns_zero(self): - lm = MagicMock() - lm.history = [{"usage": {"prompt_tokens": 100, "completion_tokens": 50}, "cost": 0.01}] - usage = usage_since(lm, 5) - assert usage.input_tokens == 0 - - def test_no_history_returns_zero(self): - lm = MagicMock(spec=[]) - usage = usage_since(lm, 0) - assert usage.input_tokens == 0 - - def test_cache_hit_via_response_flag_excluded_from_cost(self): - """DSPy's Cache.get() zeros response.usage but keeps response._hidden_params - on a cache hit. Without care, usage_since double-counts: 0 tokens yet - phantom cost. We detect cache hits via response.cache_hit and drop them - from the billed aggregate, surfacing the count via TokenUsage.cache_hits. - """ - fresh_resp = MagicMock() - fresh_resp.cache_hit = False - cached_resp = MagicMock() - cached_resp.cache_hit = True - - lm = MagicMock() - lm.history = [ - { - "usage": {"prompt_tokens": 1000, "completion_tokens": 50}, - "cost": 0.001, - "response": fresh_resp, - }, - { - "usage": {}, # DSPy clears usage on cache hit - "cost": 0.001, # but leaves cost - "response": cached_resp, - }, - ] - usage = usage_since(lm, 0) - # Only the fresh call counts toward tokens and cost. - assert usage.input_tokens == 1000 - assert usage.output_tokens == 50 - assert usage.cost == pytest.approx(0.001) - # The cache hit is surfaced for observability. - assert usage.cache_hits == 1 - - def test_cache_hit_via_empty_usage_heuristic(self): - """When the response object isn't accessible (e.g. synthetic history or - the response was dropped), an empty usage dict paired with non-zero - cost is an unambiguous cache-hit signature: a real call always - populates prompt_tokens/completion_tokens. - """ - lm = MagicMock() - lm.history = [ - {"usage": {"prompt_tokens": 1000, "completion_tokens": 50}, "cost": 0.001}, - {"usage": {}, "cost": 0.001}, # no response key, usage empty, cost > 0 - ] - usage = usage_since(lm, 0) - assert usage.input_tokens == 1000 - assert usage.output_tokens == 50 - assert usage.cost == pytest.approx(0.001) - assert usage.cache_hits == 1 - - def test_legitimate_zero_usage_entry_not_flagged(self): - """An entry with zero usage AND zero cost is a legitimate no-op, not a - cache hit. Don't surface it as one. - """ - lm = MagicMock() - lm.history = [ - {"usage": {"prompt_tokens": 0, "completion_tokens": 0}, "cost": 0}, - ] - usage = usage_since(lm, 0) - assert usage.input_tokens == 0 - assert usage.output_tokens == 0 - assert usage.cost == 0 - assert usage.cache_hits == 0 - - def test_usage_since_keeps_token_usage_free_of_truncation_metadata(self): - lm = MagicMock() - lm.history = [ - { - "usage": {"prompt_tokens": 100, "completion_tokens": 50000}, - "kwargs": {"max_tokens": 50000}, - "response": {"choices": [{"finish_reason": "length"}]}, - } - ] - - usage = usage_since(lm, 0) - - assert usage.input_tokens == 100 - assert usage.output_tokens == 50000 - assert not hasattr(usage, "truncation") - - def test_lm_finish_since_extracts_finish_reason_only(self): - lm = MagicMock() - lm.history = [ - { - "usage": {"prompt_tokens": 100, "completion_tokens": 49100}, - "kwargs": {"max_tokens": 50000}, - "response": {"choices": [{"finish_reason": "stop"}]}, - } - ] - - metadata = lm_finish_since(lm, 0) - - assert metadata == LMFinishMetadata(finish_reason="stop") + assert usage.cost == pytest.approx(0.03) + assert usage.cache_hits == 2 def test_lm_completion_metadata_includes_prompt_cache_stats(self): - lm = MagicMock() + lm = SimpleNamespace() lm.history = [ { "usage": { @@ -918,111 +450,4 @@ def test_lm_completion_metadata_includes_prompt_cache_stats(self): assert metadata.input_tokens == 1500 assert metadata.cached_input_tokens == 850 assert metadata.cache_read_ratio == pytest.approx(850 / 1500) - - -class TestConcurrentUsageAccounting: - """Regression test for the concurrency overcount bug. - - The naïve ``usage_since(shared_lm, snapshot)`` pattern inflated each - worker's delta under concurrent execution: two workers starting at the - same ``lm.history`` length and finishing after each other's entries - have landed each see the OTHER's entries in their "delta", doubling - the logged total. - - The fix is architectural: each PredictRLM instance makes its own - ``lm.copy()`` (fresh history, shared cache/callbacks/config) in - ``__init__``. Each worker's ``lm.history`` is isolated, so - ``usage_since`` sees only that worker's own calls. - """ - - def test_lm_history_delta_overcounts_under_shared_history(self): - """Demonstrates the bug shape. Two workers sharing an lm.history - both see the full delta; sum is 2x the real tokens. This is why - PredictRLM copies the lm in __init__. - """ - class SharedLM: - def __init__(self): - self.history = [] - - lm = SharedLM() - snap_A = snapshot_lm_history_len(lm) - snap_B = snapshot_lm_history_len(lm) - lm.history.append({"usage": {"prompt_tokens": 100, "completion_tokens": 10}, "cost": 0.001}) - lm.history.append({"usage": {"prompt_tokens": 200, "completion_tokens": 20}, "cost": 0.002}) - lm.history.append({"usage": {"prompt_tokens": 150, "completion_tokens": 15}, "cost": 0.0015}) - usage_A = usage_since(lm, snap_A) - usage_B = usage_since(lm, snap_B) - - real_total = sum(e["cost"] for e in lm.history) - logged_total = usage_A.cost + usage_B.cost - assert logged_total == pytest.approx(real_total * 2), ( - "shared history always inflates — that's the bug we fix via " - "per-RLM lm.copy()" - ) - - def test_lm_copy_gives_fresh_history_per_instance(self): - """The fix: dspy.LM.copy() creates a new LM with an empty - history list. Calls made through one copy don't pollute the - other's history — so two concurrent PredictRLM instances sharing - an 'original' LM still see only their own calls via usage_since. - """ - import dspy - - original = dspy.LM(model="openai/gpt-4o", cache=False) - # PredictRLM.__init__ does this internally: - lm_a = original.copy() - lm_b = original.copy() - - # Simulate calls landing in each instance's history - lm_a.history.append({"usage": {"prompt_tokens": 100, "completion_tokens": 10}, "cost": 0.001}) - lm_a.history.append({"usage": {"prompt_tokens": 150, "completion_tokens": 15}, "cost": 0.0015}) - lm_b.history.append({"usage": {"prompt_tokens": 200, "completion_tokens": 20}, "cost": 0.002}) - - # Each copy's usage_since sees only ITS own calls - u_a = usage_since(lm_a, 0) - u_b = usage_since(lm_b, 0) - assert u_a.input_tokens == 250 - assert u_b.input_tokens == 200 - # Sum matches real total — no inflation - assert u_a.cost + u_b.cost == pytest.approx(0.0045) - # And the original lm's history is untouched - assert len(original.history) == 0 - - - -class TestSanitizeForTrace: - def test_replaces_data_uri(self): - data_uri = "data:image/png;base64," + "A" * 40000 - result = _sanitize_for_trace(data_uri) - assert result == "data:image/png;base64," - - def test_replaces_nested_in_dict(self): - data = { - "page": "data:image/jpeg;base64," + "B" * 20000, - "question": "What is this?", - } - result = _sanitize_for_trace(data) - assert result["question"] == "What is this?" - assert result["page"] == "data:image/jpeg;base64," - - def test_replaces_in_list(self): - data = ["data:image/png;base64," + "C" * 10000, "normal string"] - result = _sanitize_for_trace(data) - assert result[0] == "data:image/png;base64," - assert result[1] == "normal string" - - def test_leaves_normal_strings(self): - assert _sanitize_for_trace("hello") == "hello" - assert _sanitize_for_trace("data:not-base64") == "data:not-base64" - - def test_leaves_non_strings(self): - assert _sanitize_for_trace(42) == 42 - assert _sanitize_for_trace(None) is None - - -class TestMsSince: - def test_returns_positive_int(self): - start = time.perf_counter() - ms = ms_since(start) - assert isinstance(ms, int) - assert ms >= 0 + assert lm_finish_since(lm, 0) == LMFinishMetadata(finish_reason="stop") diff --git a/tests/test_trace_on_cancellation.py b/tests/test_trace_on_cancellation.py index ab96ccab..3d0fff4d 100644 --- a/tests/test_trace_on_cancellation.py +++ b/tests/test_trace_on_cancellation.py @@ -1,199 +1,81 @@ -"""Cancellation and timeout paths must still leave a partial RunTrace. - -Background: - ``PredictRLM._forward_traced`` and ``_aforward_traced`` attach - ``exc.trace`` to any exception so the caller can diagnose partial - runs (which iteration reached, what tokens were spent, etc). The - original handler only caught ``Exception``, which meant - ``asyncio.CancelledError`` (a ``BaseException`` since 3.8) slipped - past unaugmented. In practice this lost cost accounting for every - case that hit ``asyncio.wait_for`` — ~0.5% of evaluate rollouts in - a typical long run. - - The fix widens the handler to ``except BaseException`` and wraps - ``_build_run_trace`` in a safety net so a trace-building failure - during cancellation can never mask the cancellation itself. -""" - -from __future__ import annotations +"""Cancellation preserves completed and pending work without masking the cause.""" import asyncio -import time -from types import SimpleNamespace -from unittest.mock import MagicMock +from contextlib import contextmanager +from unittest.mock import AsyncMock, MagicMock, patch +import dspy import pytest -from dspy.primitives.repl_types import REPLEntry - - -def _patch_build_run_trace(predictor, sentinel): - """Make _build_run_trace return a sentinel object we can recognise.""" - predictor._build_run_trace = MagicMock(return_value=sentinel) - - -class TestSyncForwardTracedCancellationPath: - """The sync path should attach a trace even for KeyboardInterrupt. - - We test KeyboardInterrupt rather than CancelledError because - CancelledError is an async primitive; KeyboardInterrupt is a sync - BaseException that exercises the same handler widening. - """ - - def test_keyboard_interrupt_still_attaches_trace(self): - from predict_rlm.predict_rlm import PredictRLM - - predictor = PredictRLM.__new__(PredictRLM) - sentinel = object() - _patch_build_run_trace(predictor, sentinel) - - # Build the minimum frame state _forward_traced expects so we - # can call the except branch directly by raising into it. - class _FakeExc(KeyboardInterrupt): - pass - - exc = _FakeExc() - try: - # Simulate the handler body from _forward_traced (lines - # 1426-1436): widened except sets exc.trace and re-raises. - try: - raise exc - except BaseException as e: - try: - e.trace = predictor._build_run_trace( - status="error", - steps=[], - lm=None, - sub_lm=None, - lm_hist_start=0, - sub_hist_start=0, - run_start=0, - ) - except Exception: - pass - raise - except KeyboardInterrupt as caught: - assert caught.trace is sentinel - - -class TestAsyncForwardTracedCancellationPath: - """asyncio.CancelledError must carry a trace after the widened handler.""" - - def test_cancelled_error_gets_trace_attached(self): - from predict_rlm.predict_rlm import PredictRLM - - predictor = PredictRLM.__new__(PredictRLM) - sentinel = object() - _patch_build_run_trace(predictor, sentinel) - - async def _raises_cancelled(): - try: - raise asyncio.CancelledError() - except BaseException as e: - try: - e.trace = predictor._build_run_trace( - status="error", - steps=[], - lm=None, - sub_lm=None, - lm_hist_start=0, - sub_hist_start=0, - run_start=0, - ) - except Exception: - pass - raise - - with pytest.raises(asyncio.CancelledError) as exc_info: - asyncio.run(_raises_cancelled()) - - assert exc_info.value.trace is sentinel - - def test_build_run_trace_failure_does_not_mask_cancellation(self): - """If _build_run_trace itself raises during cancellation, the - cancellation must still propagate cleanly — the ``try/except`` - around trace attachment is there exactly for this case. - """ - from predict_rlm.predict_rlm import PredictRLM - - predictor = PredictRLM.__new__(PredictRLM) - predictor._build_run_trace = MagicMock( - side_effect=RuntimeError("trace build exploded mid-cancel") - ) - - async def _raises_cancelled_with_broken_trace(): - try: - raise asyncio.CancelledError("inner") - except BaseException as e: - try: - e.trace = predictor._build_run_trace( - status="error", - steps=[], - lm=None, - sub_lm=None, - lm_hist_start=0, - sub_hist_start=0, - run_start=0, - ) - except Exception: - pass - raise - - with pytest.raises(asyncio.CancelledError): - asyncio.run(_raises_cancelled_with_broken_trace()) - - def test_error_trace_includes_pending_iteration(self): - from predict_rlm.predict_rlm import PredictRLM - - predictor = PredictRLM.__new__(PredictRLM) - predictor.max_iterations = 3 - predictor._partial_pending_entry = REPLEntry( - reasoning="about to call a slow tool", - code="await slow_tool()", - output="", - ) - predictor._partial_pending_start = time.perf_counter() - lm = SimpleNamespace(model="fake-main", history=[]) - - trace = predictor._build_run_trace( - status="error", - steps=[], - lm=lm, - sub_lm=None, - lm_hist_start=0, - sub_hist_start=0, - run_start=time.perf_counter(), - ) - - assert trace.iterations == 1 - assert trace.steps[0].reasoning == "about to call a slow tool" - assert trace.steps[0].code == "await slow_tool()" - assert trace.steps[0].error is True - - -class TestHandlerWideningIsAnchored: - """Source anchor: both error handlers in predict_rlm.py must use - ``except BaseException``, not ``except Exception``. If a future - refactor narrows them, this test fails to flag the regression. - """ - - def test_forward_traced_catches_base_exception(self): - import inspect - - from predict_rlm.predict_rlm import PredictRLM - - src = inspect.getsource(PredictRLM._forward_traced) - assert "except BaseException as exc" in src, ( - "PredictRLM._forward_traced should catch BaseException so " - "cancellations attach a partial RunTrace before re-raising" - ) - - def test_aforward_traced_catches_base_exception(self): - import inspect - - from predict_rlm.predict_rlm import PredictRLM - - src = inspect.getsource(PredictRLM._aforward_traced) - assert "except BaseException as exc" in src, ( - "PredictRLM._aforward_traced should catch BaseException so " - "asyncio.CancelledError attaches a partial RunTrace" - ) + +from predict_rlm import PredictRLM + + +def _cancelling_run(error): + class Repl: + def execute(self, code, variables=None, timeout=None): + if code == "await slow_tool()": + raise error + return "committed output" + + async def aexecute(self, code, variables=None, timeout=None): + return self.execute(code, variables, timeout) + + @contextmanager + def session(**kwargs): + yield Repl() + + rlm = PredictRLM("query -> answer", max_iterations=3) + actions = [ + dspy.Prediction(reasoning="first", code="print('committed output')"), + dspy.Prediction(reasoning="pending", code="await slow_tool()"), + ] + rlm.generate_action = MagicMock(side_effect=actions) + rlm.generate_action.acall = AsyncMock(side_effect=actions) + rlm._interpreter_context = session + return rlm + + +def _assert_partial_trace(error): + assert error.trace.status == "error" + assert [step.code for step in error.trace.steps] == [ + "print('committed output')", + "await slow_tool()", + ] + assert error.trace.steps[0].untruncated_output == "committed output" + assert error.trace.steps[1].error is True + + +def test_keyboard_interrupt_preserves_completed_and_pending_steps(): + error = KeyboardInterrupt("stop") + rlm = _cancelling_run(error) + with pytest.raises(KeyboardInterrupt) as caught: + rlm._forward_traced(None, query="work") + assert caught.value is error + _assert_partial_trace(error) + + +@pytest.mark.asyncio +async def test_async_cancellation_preserves_completed_and_pending_steps(): + error = asyncio.CancelledError("stop") + rlm = _cancelling_run(error) + with pytest.raises(asyncio.CancelledError) as caught: + await rlm._aforward_traced(None, query="work") + assert caught.value is error + _assert_partial_trace(error) + + +@pytest.mark.asyncio +async def test_trace_build_failure_does_not_replace_cancellation(): + error = asyncio.CancelledError("stop") + rlm = _cancelling_run(error) + build_trace = rlm._build_run_trace + + def fail_error_trace(*args, **kwargs): + if kwargs.get("status") == "error": + raise RuntimeError("trace unavailable") + return build_trace(*args, **kwargs) + + with patch.object(rlm, "_build_run_trace", side_effect=fail_error_trace): + with pytest.raises(asyncio.CancelledError) as caught: + await rlm._aforward_traced(None, query="work") + assert caught.value is error diff --git a/tests/test_triple_quote_preprocess.py b/tests/test_triple_quote_preprocess.py deleted file mode 100644 index 7e6cb1b3..00000000 --- a/tests/test_triple_quote_preprocess.py +++ /dev/null @@ -1,87 +0,0 @@ -"""Tests for the nested triple-quote SyntaxError handling. - -When an RLM emits Python code that wraps a multi-line string in ``\"\"\"`` -and that content itself contains ``\"\"\"``, Python's tokenizer fails. We -do NOT auto-rewrite the code — the agent gets an enhanced error message -with an actionable hint and can self-correct on its next REPL iteration. - -These tests anchor the Python-level bug (so we notice if Python ever -fixes it) and verify the detection heuristic used by the enhanced error -path in ``PredictRLM._aexecute_iteration``. -""" - -from __future__ import annotations - -import ast - -import pytest - - -class TestPythonParserBugAnchor: - """Anchor: confirm Python itself rejects nested triple-quotes. - - If these ever stop failing, Python changed its parser and the - enhanced error hint in PredictRLM is no longer needed. - """ - - def test_nested_triple_quote_literal_rejected_by_parser(self): - broken = 'x = """pre\n"""inside"""\npost"""\nprint(x)' - with pytest.raises(SyntaxError): - ast.parse(broken) - - def test_nested_triple_quote_with_code_block_content_rejected(self): - broken = ( - 'new_instructions = """# Skill\n' - "\n" - "```python\n" - "def compute():\n" - ' """docstring lives here"""\n' - " return 42\n" - "```\n" - '"""\n' - "SUBMIT(new_instructions=new_instructions)\n" - ) - with pytest.raises(SyntaxError): - ast.parse(broken) - - -class TestTripleQuoteDetectionHeuristic: - """The detection heuristic: code.count('\"\"\"') >= 3 plus error keywords.""" - - def test_heuristic_fires_on_even_count(self): - code = 'x = """pre\n"""inside"""\npost"""\nprint(x)' - assert code.count('"""') >= 3 - with pytest.raises(SyntaxError): - ast.parse(code) - err = "" - try: - ast.parse(code) - except SyntaxError as e: - err = str(e).lower() - assert "invalid syntax" in err or "unterminated" in err - - def test_heuristic_fires_on_odd_count(self): - code = ( - 'new_instructions = """# Skill\n' - "\n" - "```python\n" - "def compute():\n" - ' """docstring lives here"""\n' - " return 42\n" - "```\n" - '"""\n' - ) - assert code.count('"""') >= 3 - with pytest.raises(SyntaxError): - ast.parse(code) - - def test_heuristic_does_not_fire_on_valid_triple_quoted_string(self): - code = 'x = """hello\nworld"""\nprint(x)' - assert code.count('"""') == 2 # exactly 2 = one pair, no nesting - ast.parse(code) # should not raise - - def test_heuristic_does_not_fire_on_non_triple_quote_syntax_error(self): - code = "x = \n" - assert code.count('"""') == 0 - with pytest.raises(SyntaxError): - ast.parse(code) diff --git a/tests/test_workspace.py b/tests/test_workspace.py index a0f437bc..22382db7 100644 --- a/tests/test_workspace.py +++ b/tests/test_workspace.py @@ -2,9 +2,7 @@ from __future__ import annotations -import asyncio import hashlib -import json import os import shutil import sys @@ -15,7 +13,7 @@ import dspy import pytest -from predict_rlm import PredictRLM, Workspace, WorkspaceMode +from predict_rlm import PredictRLM, Workspace from predict_rlm.compatibility import WorkspaceInputAdapter from predict_rlm.runtime import ( ArtifactBinding, @@ -23,9 +21,6 @@ ExecutionResult, FieldDescriptor, FileTransfer, - HostDirectoryMount, - PreparedInput, - UnsupportedOperationError, ) from predict_rlm.workspace import ( WorkspaceFileInfo, @@ -43,28 +38,6 @@ ) -class TestWorkspace: - def test_exported_from_package(self): - from predict_rlm import Workspace as ExportedWorkspace - - assert ExportedWorkspace is Workspace - - def test_create_with_defaults(self): - workspace = Workspace(path="/tmp/repo") - assert workspace.path == "/tmp/repo" - assert workspace.mount_path == "/sandbox/workspace" - assert workspace.mode is WorkspaceMode.MIRROR - assert workspace.sync_back is True - assert ".git" in workspace.exclude - assert "node_modules" in workspace.exclude - assert workspace.max_file_bytes == 5_000_000 - - def test_create_direct_mode(self): - workspace = Workspace(path="/tmp/repo", mount_path="/workspace", mode="direct") - assert workspace.mode is WorkspaceMode.DIRECT - assert workspace.mount_path == "/workspace" - - class TestWorkspaceSyncState: def _info(self, text: str) -> WorkspaceFileInfo: data = text.encode() @@ -89,7 +62,6 @@ def test_skipped_large_host_file_conflicts_instead_of_clobbering(self): with pytest.raises(WorkspaceSyncConflictError, match="large.txt"): state.sync_from_sandbox(repl) - repl.sync_file_to.assert_not_called() with open(host_path) as f: assert f.read() == "too large" @@ -108,7 +80,6 @@ def test_oversized_sandbox_rewrite_conflicts_instead_of_deleting_host_file(self) with pytest.raises(WorkspaceSyncConflictError, match="small.txt"): state.sync_from_sandbox(repl) - repl.sync_file_to.assert_not_called() with open(host_path) as f: assert f.read() == "small" @@ -131,7 +102,6 @@ def test_host_symlink_skipped_from_manifest_and_conflicts_on_write(self): with pytest.raises(WorkspaceSyncConflictError, match="link.txt"): state.sync_from_sandbox(repl) - repl.sync_file_to.assert_not_called() with open(target) as f: assert f.read() == "target" @@ -168,7 +138,6 @@ def test_conflict_detection_is_atomic_before_any_write(self): with pytest.raises(WorkspaceSyncConflictError, match="conflict.txt"): state.sync_from_sandbox(repl) - repl.sync_file_to.assert_not_called() with open(clean) as f: assert f.read() == "base clean" @@ -187,7 +156,6 @@ def test_workspace_manifest_failure_conflicts_instead_of_deleting_host_files(sel with pytest.raises(WorkspaceSyncConflictError, match="mount disappeared"): state.sync_from_sandbox(repl) - repl.sync_file_to.assert_not_called() with open(path) as f: assert f.read() == "keep" @@ -204,274 +172,40 @@ def bind(self, binding: ArtifactBinding) -> None: class _WorkspaceTransportSession: name = "workspace-transport" - def __init__(self) -> None: - self.direct_mounts: list[HostDirectoryMount] = [] - self.direct_mount_path: str | None = None - self.created_directories: list[str] = [] - self.transfers: list[FileTransfer] = [] - self.sandbox_files: dict[str, bytes] = {} - self.inspect_calls: list[str] = [] - self.collect_calls: list[tuple[str, str]] = [] + def __init__(self, root: Path) -> None: + self.root = root + + def path(self, sandbox_path: str) -> Path: + return self.root / sandbox_path.lstrip("/") async def transfer_file(self, transfer: FileTransfer) -> str: - self.transfers.append(transfer) - self.sandbox_files[transfer.sandbox_path] = Path( - transfer.source_path - ).read_bytes() + target = self.path(transfer.sandbox_path) + target.parent.mkdir(parents=True, exist_ok=True) + shutil.copy2(transfer.source_path, target) return transfer.sandbox_path - async def mount_host_directory(self, mount: HostDirectoryMount) -> str: - self.direct_mounts.append(mount) - return self.direct_mount_path or mount.sandbox_path - async def create_directory(self, sandbox_path: str) -> None: - self.created_directories.append(sandbox_path) + self.path(sandbox_path).mkdir(parents=True, exist_ok=True) async def inspect_directory(self, sandbox_path: str): - self.inspect_calls.append(sandbox_path) - prefix = sandbox_path.rstrip("/") + "/" + root = self.path(sandbox_path) return { - path.removeprefix(prefix): ArtifactFileInfo( + path.relative_to(root).as_posix(): ArtifactFileInfo( type="file", - sha256=hashlib.sha256(contents).hexdigest(), - size=len(contents), + sha256=hashlib.sha256(path.read_bytes()).hexdigest(), + size=path.stat().st_size, ) - for path, contents in self.sandbox_files.items() - if path.startswith(prefix) + for path in root.rglob("*") + if path.is_file() } async def collect_file(self, sandbox_path: str, host_path: str) -> None: - self.collect_calls.append((sandbox_path, host_path)) destination = Path(host_path) destination.parent.mkdir(parents=True, exist_ok=True) - destination.write_bytes(self.sandbox_files[sandbox_path]) - - -class _CopyInOnlyWorkspaceSession: - name = "copy-in-only" - - def __init__(self) -> None: - self.created_directories: list[str] = [] - self.transfers: list[FileTransfer] = [] - - async def create_directory(self, sandbox_path: str) -> None: - self.created_directories.append(sandbox_path) - - async def transfer_file(self, transfer: FileTransfer) -> str: - self.transfers.append(transfer) - return transfer.sandbox_path + shutil.copy2(self.path(sandbox_path), destination) class TestWorkspaceInputAdapterLifecycle: - @pytest.mark.asyncio - async def test_workspace_lifecycle_rejects_prepared_input_without_typed_state(self): - with pytest.raises(TypeError, match="typed Workspace prepared state"): - await WorkspaceInputAdapter().bind( - FieldDescriptor("workspace", Workspace), - PreparedInput(model_value="/sandbox/workspace"), - _WorkspaceAdapterContext(), - _WorkspaceTransportSession(), - ) - - @pytest.mark.asyncio - async def test_copy_in_only_workspace_does_not_require_sync_back_capabilities( - self, - tmp_path: Path, - ): - workspace_root = tmp_path / "workspace" - workspace_root.mkdir() - (workspace_root / "source.txt").write_text("before", encoding="utf-8") - field = FieldDescriptor("workspace", Workspace) - ctx = _WorkspaceAdapterContext() - adapter = WorkspaceInputAdapter() - prepared = await adapter.prepare( - field, - Workspace(path=str(workspace_root), sync_back=False), - ctx, - ) - session = _CopyInOnlyWorkspaceSession() - - bound = await adapter.bind(field, prepared, ctx, session) - await adapter.finalize(field, prepared, ctx, session, None) - - assert bound.model_value == "/sandbox/workspace" - assert session.created_directories == ["/sandbox/workspace"] - assert [transfer.sandbox_path for transfer in session.transfers] == [ - "/sandbox/workspace/source.txt" - ] - - @pytest.mark.asyncio - async def test_prepare_uses_neutral_artifacts_and_typed_requirements( - self, - tmp_path: Path, - ): - workspace_root = tmp_path / "workspace" - workspace_root.mkdir() - workspace = Workspace( - path=str(workspace_root), - mount_path="/workspace", - mode=WorkspaceMode.DIRECT, - ) - field = FieldDescriptor("workspace", Workspace) - ctx = _WorkspaceAdapterContext() - adapter = WorkspaceInputAdapter() - - prepared = await adapter.prepare(field, workspace, ctx) - expected_mount = HostDirectoryMount( - host_path=str(workspace_root.resolve()), - sandbox_path="/workspace", - ) - assert prepared.model_value == "/workspace" - assert len(prepared.artifacts) == 1 - artifact = prepared.artifacts[0] - assert artifact.kind == "compat.workspace" - assert dict(artifact.metadata) == {"sandbox_path": "/workspace"} - assert "workspace_binding" not in artifact.metadata - assert "compat.workspace.direct" not in artifact.kind - assert "compat.workspace.mirror" not in artifact.kind - json.dumps(dict(artifact.metadata)) - assert prepared.host_directory_mounts == (expected_mount,) - assert prepared.requirements.extra_read_paths == (str(workspace_root.resolve()),) - assert prepared.requirements.extra_write_paths == (str(workspace_root.resolve()),) - - @pytest.mark.asyncio - async def test_direct_mount_uses_exact_host_directory_primitive(self, tmp_path: Path): - workspace_root = tmp_path / "workspace" - workspace_root.mkdir() - field = FieldDescriptor("workspace", Workspace) - ctx = _WorkspaceAdapterContext() - adapter = WorkspaceInputAdapter() - prepared = await adapter.prepare( - field, - Workspace( - path=str(workspace_root), - mount_path="/workspace", - mode=WorkspaceMode.DIRECT, - ), - ctx, - ) - session = _WorkspaceTransportSession() - session.direct_mount_path = "/mounted/workspace" - - bound = await adapter.bind(field, prepared, ctx, session) - - assert session.direct_mounts == [ - HostDirectoryMount( - host_path=str(workspace_root.resolve()), - sandbox_path="/workspace", - ) - ] - assert session.transfers == [] - assert bound.model_value == "/mounted/workspace" - assert [binding.path for binding in bound.bindings] == [ - "/mounted/workspace" - ] - - @pytest.mark.asyncio - async def test_direct_mount_is_rejected_by_jspi(self, tmp_path: Path): - from predict_rlm.backends.jspi import JspiExecutionBackend - - mount = HostDirectoryMount(str(tmp_path), "/workspace") - ctx = MagicMock(spec=None) - - with pytest.raises(UnsupportedOperationError, match="JSPI"): - await JspiExecutionBackend().validate_host_directory_mounts( - (mount,), - ctx, - ) - - @pytest.mark.sbx - @pytest.mark.asyncio - async def test_direct_mount_is_rejected_by_pooled_sbx(self, tmp_path: Path): - from predict_rlm.backends.sbx import SbxPoolExecutionBackend - - mount = HostDirectoryMount(str(tmp_path), "/workspace") - ctx = MagicMock(spec=None) - - with pytest.raises(UnsupportedOperationError, match="SbxPool"): - await SbxPoolExecutionBackend(MagicMock()).validate_host_directory_mounts( - (mount,), - ctx, - ) - - @pytest.mark.asyncio - async def test_mirror_bind_and_sync_use_generic_transport_after_success_and_failure( - self, - tmp_path: Path, - ): - workspace_root = tmp_path / "workspace" - workspace_root.mkdir() - original = workspace_root / "original.txt" - deleted = workspace_root / "deleted.txt" - original.write_text("before", encoding="utf-8") - deleted.write_text("delete", encoding="utf-8") - field = FieldDescriptor("workspace", Workspace) - ctx = _WorkspaceAdapterContext() - adapter = WorkspaceInputAdapter() - prepared = await adapter.prepare(field, Workspace(path=str(workspace_root)), ctx) - session = _WorkspaceTransportSession() - - bound = await adapter.bind(field, prepared, ctx, session) - - assert session.created_directories == ["/sandbox/workspace"] - assert {transfer.sandbox_path for transfer in session.transfers} == { - "/sandbox/workspace/deleted.txt", - "/sandbox/workspace/original.txt", - } - assert bound.model_value == "/sandbox/workspace" - - session.sandbox_files["/sandbox/workspace/original.txt"] = b"after success" - session.sandbox_files["/sandbox/workspace/created.txt"] = b"created" - del session.sandbox_files["/sandbox/workspace/deleted.txt"] - await adapter.after_execution( - field, - prepared, - ctx, - session, - ExecutionResult(value="ok"), - None, - ) - - assert original.read_text(encoding="utf-8") == "after success" - assert (workspace_root / "created.txt").read_text(encoding="utf-8") == "created" - assert not deleted.exists() - - session.sandbox_files["/sandbox/workspace/original.txt"] = b"after failure" - execution_error = RuntimeError("generated code failed") - await adapter.after_execution( - field, - prepared, - ctx, - session, - None, - execution_error, - ) - - assert original.read_text(encoding="utf-8") == "after failure" - assert session.inspect_calls == ["/sandbox/workspace", "/sandbox/workspace"] - - @pytest.mark.asyncio - async def test_finalize_syncs_once_and_is_idempotent(self, tmp_path: Path): - workspace_root = tmp_path / "workspace" - workspace_root.mkdir() - source = workspace_root / "source.txt" - source.write_text("before", encoding="utf-8") - field = FieldDescriptor("workspace", Workspace) - ctx = _WorkspaceAdapterContext() - adapter = WorkspaceInputAdapter() - prepared = await adapter.prepare(field, Workspace(path=str(workspace_root)), ctx) - session = _WorkspaceTransportSession() - await adapter.bind(field, prepared, ctx, session) - session.sandbox_files["/sandbox/workspace/source.txt"] = b"final" - await adapter.finalize(field, prepared, ctx, session, None) - await adapter.finalize(field, prepared, ctx, session, None) - - assert source.read_text(encoding="utf-8") == "final" - assert session.inspect_calls == ["/sandbox/workspace"] - assert session.collect_calls == [ - ("/sandbox/workspace/source.txt", str(source.resolve())) - ] - @pytest.mark.asyncio async def test_finalize_continues_after_list_item_conflict(self, tmp_path: Path): first_root = tmp_path / "first" @@ -493,11 +227,11 @@ async def test_finalize_continues_after_list_item_conflict(self, tmp_path: Path) ], ctx, ) - session = _WorkspaceTransportSession() + session = _WorkspaceTransportSession(tmp_path / "sandbox") await adapter.bind(field, prepared, ctx, session) first_source.write_text("host concurrent change", encoding="utf-8") - session.sandbox_files["/sandbox/first/source.txt"] = b"sandbox change" - session.sandbox_files["/sandbox/second/source.txt"] = b"synced" + session.path("/sandbox/first/source.txt").write_text("sandbox change") + session.path("/sandbox/second/source.txt").write_text("synced") with pytest.raises(WorkspaceSyncConflictError, match="source.txt"): await adapter.after_execution( field, @@ -514,71 +248,6 @@ async def test_finalize_continues_after_list_item_conflict(self, tmp_path: Path) assert first_source.read_text(encoding="utf-8") == "host concurrent change" assert second_source.read_text(encoding="utf-8") == "synced" - @pytest.mark.asyncio - async def test_finalize_flushes_mirror_after_cancelled_execution(self, tmp_path: Path): - workspace_root = tmp_path / "workspace" - workspace_root.mkdir() - source = workspace_root / "source.txt" - source.write_text("before", encoding="utf-8") - field = FieldDescriptor("workspace", Workspace) - ctx = _WorkspaceAdapterContext() - adapter = WorkspaceInputAdapter() - prepared = await adapter.prepare(field, Workspace(path=str(workspace_root)), ctx) - session = _WorkspaceTransportSession() - await adapter.bind(field, prepared, ctx, session) - session.sandbox_files["/sandbox/workspace/source.txt"] = b"after cancellation" - - await adapter.finalize( - field, - prepared, - ctx, - session, - asyncio.CancelledError(), - ) - - assert source.read_text(encoding="utf-8") == "after cancellation" - assert session.inspect_calls == ["/sandbox/workspace"] - - @pytest.mark.asyncio - async def test_mirror_conflict_through_generic_transport_preserves_host_change( - self, - tmp_path: Path, - ): - workspace_root = tmp_path / "workspace" - workspace_root.mkdir() - source = workspace_root / "source.txt" - source.write_text("base", encoding="utf-8") - field = FieldDescriptor("workspace", Workspace) - ctx = _WorkspaceAdapterContext() - adapter = WorkspaceInputAdapter() - prepared = await adapter.prepare(field, Workspace(path=str(workspace_root)), ctx) - session = _WorkspaceTransportSession() - await adapter.bind(field, prepared, ctx, session) - source.write_text("host concurrent change", encoding="utf-8") - session.sandbox_files["/sandbox/workspace/source.txt"] = b"sandbox change" - - with pytest.raises(WorkspaceSyncConflictError, match="source.txt"): - await adapter.after_execution( - field, - prepared, - ctx, - session, - None, - RuntimeError("generated code failed"), - ) - - assert source.read_text(encoding="utf-8") == "host concurrent change" - assert session.collect_calls == [] - - @pytest.mark.asyncio - async def test_missing_workspace_fails_during_adapter_preparation(self): - with pytest.raises(FileNotFoundError, match="/no/such/workspace"): - await WorkspaceInputAdapter().prepare( - FieldDescriptor("workspace", Workspace), - Workspace(path="/no/such/workspace"), - _WorkspaceAdapterContext(), - ) - class _WorkspaceMutationActions: async def acall(self, *, iteration, **kwargs): diff --git a/tests/test_wt_hooks.py b/tests/test_wt_hooks.py deleted file mode 100644 index e46cc5f3..00000000 --- a/tests/test_wt_hooks.py +++ /dev/null @@ -1,28 +0,0 @@ -import os -import subprocess -from pathlib import Path - -import pytest - - -@pytest.mark.parametrize("worktree_name", ["predict-rlm.feature", "predict-rlm-feature"]) -def test_setup_hook_links_env_development_from_sibling_worktree(tmp_path, worktree_name): - main_repo = tmp_path / "predict-rlm" - worktree = tmp_path / worktree_name - main_repo.mkdir() - worktree.mkdir() - (main_repo / ".env.development").write_text("TOKEN=dev\n") - - hook = Path(__file__).resolve().parents[1] / ".wt" / "hooks" / "setup.sh" - - subprocess.run( - [str(hook), str(worktree), worktree_name], - check=True, - text=True, - capture_output=True, - ) - - env_link = worktree / ".env.development" - assert env_link.is_symlink() - assert os.readlink(env_link) == "../predict-rlm/.env.development" - assert env_link.resolve() == main_repo / ".env.development" From 2196e887ac28e095899d14154db898041f63b66f Mon Sep 17 00:00:00 2001 From: Emile Riberdy Date: Tue, 8 Sep 2026 14:39:27 +0000 Subject: [PATCH 2/5] test(predict-rlm): synchronize tool callbacks across worker loops --- .github/workflows/tests.yml | 7 ++++++- tests/runtime_contracts/test_tool_contract.py | 11 ++--------- 2 files changed, 8 insertions(+), 10 deletions(-) diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index 030799eb..0bb61544 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -156,7 +156,12 @@ jobs: run: uv sync --python 3.12 - name: Run JSPI integration tests - run: uv run pytest tests/ -m "integration and not sbx and not local" -q --cov=predict_rlm --cov-report=xml + env: + PYTHONUNBUFFERED: "1" + run: >- + uv run pytest tests/ -m "integration and not sbx and not local" + -o addopts= -v --tb=short -o faulthandler_timeout=60 + --cov=predict_rlm --cov-report=xml - name: Upload coverage to Codecov uses: codecov/codecov-action@v5 diff --git a/tests/runtime_contracts/test_tool_contract.py b/tests/runtime_contracts/test_tool_contract.py index 72c2552f..1fba3ee4 100644 --- a/tests/runtime_contracts/test_tool_contract.py +++ b/tests/runtime_contracts/test_tool_contract.py @@ -30,21 +30,14 @@ def test_host_tool_result_shapes(runtime: RuntimeHandle) -> None: def test_host_tools_run_concurrently(runtime: RuntimeHandle, asynchronous: bool) -> None: runtime.require("concurrent_tools") barrier = threading.Barrier(2) - arrivals = 0 - both_started = None def sync_tool(value): barrier.wait(timeout=2) return value * 2 async def async_tool(value): - nonlocal arrivals, both_started - if both_started is None: - both_started = asyncio.Event() - arrivals += 1 - if arrivals == 2: - both_started.set() - await asyncio.wait_for(both_started.wait(), timeout=2) + # Sync SBX dispatch gives each callback its own worker event loop. + await asyncio.to_thread(barrier.wait, timeout=2) return value * 2 runtime.configure(tools={"double": async_tool if asynchronous else sync_tool}) From 5dfc499ff5f764b30f59a2584bdbf0d548dc3b06 Mon Sep 17 00:00:00 2001 From: Emile Riberdy Date: Tue, 8 Sep 2026 14:57:58 +0000 Subject: [PATCH 3/5] fix(predict-rlm): quiesce JSPI watchdog before timeout cleanup --- CHANGELOG.md | 6 +++ src/predict_rlm/backends/jspi/payload.js | 40 +++++++++++-------- .../test_external_input_adapter_contracts.py | 12 +++++- 3 files changed, 40 insertions(+), 18 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 6aeb560a..31260c4e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -14,6 +14,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 static/presentation checks, and one-off benchmark/example checks. - Direct CPython contracts now run in the core tier without the SBX extra. +### Fixed + +- JSPI timeout cleanup now waits for the interrupt worker to disarm before + entering Python, and preserves the timeout response if tracing interrupts + cleanup instead of leaving the host waiting for an uncorrelated error. + ### Breaking Changes - `SbxBackend.shutdown()` can no longer be called from the event loop that owns diff --git a/src/predict_rlm/backends/jspi/payload.js b/src/predict_rlm/backends/jspi/payload.js index d3638e70..00730478 100644 --- a/src/predict_rlm/backends/jspi/payload.js +++ b/src/predict_rlm/backends/jspi/payload.js @@ -532,7 +532,7 @@ const SIGINT = 2; const interruptBuffer = new Int32Array(new SharedArrayBuffer(8)); let interruptTimerWorker = null; let interruptArmId = 0; -const pendingInterruptArms = new Map(); +const pendingInterruptRequests = new Map(); try { pyodide.setInterruptBuffer(interruptBuffer); @@ -581,6 +581,7 @@ self.onmessage = (event) => { self.postMessage({ type: "armed", armId: data.armId }); } else if (data.type === "disarm") { clearTimer(true); + self.postMessage({ type: "disarmed", armId: data.armId }); } }; `], { type: "application/javascript" })), { type: "module" }); @@ -593,11 +594,11 @@ self.onmessage = (event) => { resolve(data.cancellationSignalReady === true); return; } - if (data.type !== "armed") return; - const pending = pendingInterruptArms.get(data.armId); + if (data.type !== "armed" && data.type !== "disarmed") return; + const pending = pendingInterruptRequests.get(data.armId); if (!pending) return; clearTimeout(pending.timerId); - pendingInterruptArms.delete(data.armId); + pendingInterruptRequests.delete(data.armId); pending.resolve(true); }; }); @@ -616,10 +617,10 @@ const armExecutionInterrupt = async (timeoutSeconds) => { const armId = ++interruptArmId; const armed = new Promise((resolve) => { const timerId = setTimeout(() => { - pendingInterruptArms.delete(armId); + pendingInterruptRequests.delete(armId); resolve(false); }, 1000); - pendingInterruptArms.set(armId, { resolve, timerId }); + pendingInterruptRequests.set(armId, { resolve, timerId }); }); interruptTimerWorker.postMessage({ type: "arm", @@ -630,14 +631,19 @@ const armExecutionInterrupt = async (timeoutSeconds) => { return await armed; }; -const disarmExecutionInterrupt = () => { - if (interruptTimerWorker) { - interruptTimerWorker.postMessage({ type: "disarm", buffer: interruptBuffer }); - } - for (const [armId, pending] of pendingInterruptArms) { +const disarmExecutionInterrupt = async () => { + for (const [armId, pending] of pendingInterruptRequests) { clearTimeout(pending.timerId); pending.resolve(false); - pendingInterruptArms.delete(armId); + pendingInterruptRequests.delete(armId); + } + if (interruptTimerWorker) { + const armId = ++interruptArmId; + const disarmed = new Promise((resolve) => { + pendingInterruptRequests.set(armId, { resolve, timerId: null }); + }); + interruptTimerWorker.postMessage({ type: "disarm", buffer: interruptBuffer, armId }); + await disarmed; } Atomics.store(interruptBuffer, 0, 0); Atomics.store(interruptBuffer, 1, 0); @@ -667,12 +673,12 @@ sys.settrace(__predict_rlm_timeout_trace) }; const disablePythonExecutionTimeout = () => { - pyodide.globals.set("__predict_rlm_timeout_disabled", true); try { + pyodide.globals.set("__predict_rlm_timeout_disabled", true); pyodide.runPython("import sys\nsys.settrace(None)"); } catch (e) { - // The JS interrupt buffer is still the fallback for states where Python - // tracing cannot safely run cleanup code. + // CPython clears a trace callback that raises while cleanup enters Python. + console.error(`[timeout] Python tracing interrupted cleanup: ${e}`); } }; @@ -1294,10 +1300,10 @@ await micropip.install([__predict_rlm_package], verbose=False) enablePythonExecutionTimeout(executionTimeoutSeconds); } const result = await pyodide.runPythonAsync(code); + await disarmExecutionInterrupt(); if (hasExecutionTimeout) { disablePythonExecutionTimeout(); } - disarmExecutionInterrupt(); // Signal code execution complete codeExecutionInProgress = false; @@ -1323,10 +1329,10 @@ await micropip.install([__predict_rlm_package], verbose=False) console.log(jsonrpcResult({ output }, requestId)); } catch (error) { const executionTimeoutFired = Atomics.load(interruptBuffer, 1) === 1; + await disarmExecutionInterrupt(); if (hasExecutionTimeout) { disablePythonExecutionTimeout(); } - disarmExecutionInterrupt(); codeExecutionInProgress = false; // Signal the response reader to stop immediately diff --git a/tests/test_external_input_adapter_contracts.py b/tests/test_external_input_adapter_contracts.py index d7544083..64427552 100644 --- a/tests/test_external_input_adapter_contracts.py +++ b/tests/test_external_input_adapter_contracts.py @@ -162,7 +162,17 @@ async def acall(self, *, iteration, **kwargs): @pytest.mark.integration @pytest.mark.skipif(shutil.which("deno") is None, reason="requires Deno") @pytest.mark.asyncio -async def test_one_stateless_adapter_handles_interleaved_real_jspi_runs(tmp_path: Path): +async def test_one_stateless_adapter_handles_interleaved_real_jspi_runs( + tmp_path: Path, monkeypatch +): + from functools import partial + + from predict_rlm.backends import JspiBackend + + monkeypatch.setattr( + "predict_rlm.backends.jspi.execution.JspiBackend", + partial(JspiBackend, preinstall_packages=False), + ) adapter = MutableRepositoryAdapter() repositories = [] for name in ("first", "second"): From a96ac237254edb711fa6fc8aedde8db0058c96bf Mon Sep 17 00:00:00 2001 From: Emile Riberdy Date: Tue, 8 Sep 2026 16:00:36 +0000 Subject: [PATCH 4/5] fix(predict-rlm): isolate deadlines from asyncio scheduler callbacks --- ARCHITECTURE.md | 7 ++- CHANGELOG.md | 3 + src/predict_rlm/backends/jspi/payload.js | 28 +++++++-- tests/runtime_contracts/test_tool_contract.py | 8 ++- tests/test_iteration_execution_timeout.py | 58 +++++++++++++++++-- 5 files changed, 89 insertions(+), 15 deletions(-) diff --git a/ARCHITECTURE.md b/ARCHITECTURE.md index d18b3a2c..768a988e 100644 --- a/ARCHITECTURE.md +++ b/ARCHITECTURE.md @@ -256,8 +256,11 @@ Python error. ### JSPI / Deno / Pyodide - Successful execute: same live Pyodide VM; full globals persist. -- Cooperative timeout: Pyodide is interrupted by trace deadline and JS interrupt - buffer; stdout/stderr are returned and the VM remains live. +- Cooperative timeout: deadline tracing is scoped to generated code. The JS + interrupt buffer stops active user code; an interrupt in an asyncio scheduler + callback cancels the suspended execution task instead of raising through the + scheduler. stdout/stderr are returned and the VM remains live. Cleanup waits + for the watchdog to disarm and restores the previous Python SIGINT handler. - Hard timeout or crash: the host watchdog kills Deno and raises `SandboxFatalError`; the interpreter is dead and no state is recovered. diff --git a/CHANGELOG.md b/CHANGELOG.md index 31260c4e..2d7302cd 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -19,6 +19,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - JSPI timeout cleanup now waits for the interrupt worker to disarm before entering Python, and preserves the timeout response if tracing interrupts cleanup instead of leaving the host waiting for an uncorrelated error. +- JSPI deadlines no longer interrupt asyncio scheduler callbacks; suspended + executions are cancelled without orphaning their result promises, and the + previous Python SIGINT handler is restored after bounded execution. ### Breaking Changes diff --git a/src/predict_rlm/backends/jspi/payload.js b/src/predict_rlm/backends/jspi/payload.js index 00730478..7495a480 100644 --- a/src/predict_rlm/backends/jspi/payload.js +++ b/src/predict_rlm/backends/jspi/payload.js @@ -653,6 +653,8 @@ const enablePythonExecutionTimeout = (timeoutSeconds) => { pyodide.globals.set("__predict_rlm_timeout_seconds", timeoutSeconds); pyodide.globals.set("__predict_rlm_timeout_disabled", false); pyodide.runPython(` +import asyncio +import signal import sys import time @@ -660,14 +662,32 @@ class PredictRLMExecutionTimeout(BaseException): pass __predict_rlm_timeout_deadline = time.monotonic() + float(__predict_rlm_timeout_seconds) +__predict_rlm_execution_task = None +__predict_rlm_previous_sigint = signal.getsignal(signal.SIGINT) + +def __predict_rlm_execution_interrupt(signum, frame): + while frame is not None: + if frame.f_code.co_filename == "": + if time.monotonic() >= __predict_rlm_timeout_deadline: + raise PredictRLMExecutionTimeout() + raise KeyboardInterrupt() + frame = frame.f_back + if __predict_rlm_execution_task is not None: + __predict_rlm_execution_task.cancel() def __predict_rlm_timeout_trace(frame, event, arg): + global __predict_rlm_execution_task if globals().get("__predict_rlm_timeout_disabled", False): return None + if frame.f_code.co_filename != "": + return None + if __predict_rlm_execution_task is None: + __predict_rlm_execution_task = asyncio.current_task() if time.monotonic() >= __predict_rlm_timeout_deadline: raise PredictRLMExecutionTimeout() return __predict_rlm_timeout_trace +signal.signal(signal.SIGINT, __predict_rlm_execution_interrupt) sys.settrace(__predict_rlm_timeout_trace) `); }; @@ -675,7 +695,7 @@ sys.settrace(__predict_rlm_timeout_trace) const disablePythonExecutionTimeout = () => { try { pyodide.globals.set("__predict_rlm_timeout_disabled", true); - pyodide.runPython("import sys\nsys.settrace(None)"); + pyodide.runPython("import sys, signal\nsys.settrace(None)\nsignal.signal(signal.SIGINT, __predict_rlm_previous_sigint)\n__predict_rlm_execution_task = None"); } catch (e) { // CPython clears a trace callback that raises while cleanup enters Python. console.error(`[timeout] Python tracing interrupted cleanup: ${e}`); @@ -1296,10 +1316,10 @@ await micropip.install([__predict_rlm_package], verbose=False) // Run the user's code if (hasExecutionTimeout) { - await armExecutionInterrupt(executionTimeoutSeconds); enablePythonExecutionTimeout(executionTimeoutSeconds); + await armExecutionInterrupt(executionTimeoutSeconds); } - const result = await pyodide.runPythonAsync(code); + const result = await pyodide.runPythonAsync(code, { filename: "" }); await disarmExecutionInterrupt(); if (hasExecutionTimeout) { disablePythonExecutionTimeout(); @@ -1404,7 +1424,7 @@ await micropip.install([__predict_rlm_package], verbose=False) hasExecutionTimeout && ( pythonExecutionTimeoutFired || - (executionTimeoutFired && errorType === "KeyboardInterrupt") + (executionTimeoutFired && (errorType === "KeyboardInterrupt" || errorType === "CancelledError")) ) ) { let capturedStdout = ""; diff --git a/tests/runtime_contracts/test_tool_contract.py b/tests/runtime_contracts/test_tool_contract.py index 1fba3ee4..0970db38 100644 --- a/tests/runtime_contracts/test_tool_contract.py +++ b/tests/runtime_contracts/test_tool_contract.py @@ -117,14 +117,16 @@ def test_timeout_during_concurrent_host_tools_is_recoverable( runtime: RuntimeHandle, tmp_path: Path, ) -> None: - result_queue = multiprocessing.Queue() - process = multiprocessing.Process( + context = multiprocessing.get_context("spawn") + result_queue = context.Queue() + process = context.Process( target=_run_timeout_repro, args=(runtime.spec.name, str(tmp_path / "staging"), result_queue), ) process.start() try: - status, *payload = _get_message(process, result_queue, 30) + # Spawn imports the Python dependencies before the backend's own startup. + status, *payload = _get_message(process, result_queue, 60) if status == "skip": pytest.skip(payload[0]) assert status == "ready", payload diff --git a/tests/test_iteration_execution_timeout.py b/tests/test_iteration_execution_timeout.py index 962abe8a..a4296210 100644 --- a/tests/test_iteration_execution_timeout.py +++ b/tests/test_iteration_execution_timeout.py @@ -82,17 +82,30 @@ async def _silent_execute(_request_id): @pytest.mark.integration -def test_predict_rlm_jspi_timeout_preserves_state_history_and_predict_tool(): +def test_predict_rlm_jspi_timeout_preserves_state_history_and_predict_tool(monkeypatch): + from functools import partial + from predict_rlm import PredictRLM + from predict_rlm.backends import JspiBackend from predict_rlm.predict_rlm import dspy + monkeypatch.setattr( + "predict_rlm.backends.jspi.execution.JspiBackend", + partial(JspiBackend, preinstall_packages=False), + ) + actions = _SequentialActions( SimpleNamespace( - reasoning="call predict before a bounded risky loop", + reasoning="prepare state before the bounded operation", code=( "first = await predict('question: str -> answer: str', " "question='first call')\n" "saved = {'first': first['answer'], 'marker': 123}\n" + ), + ), + SimpleNamespace( + reasoning="run a bounded risky loop using prepared state", + code=( "print('first predict:', saved['first'])\n" "print('marker before timeout:', saved['marker'])\n" "while True:\n" @@ -121,7 +134,7 @@ def test_predict_rlm_jspi_timeout_preserves_state_history_and_predict_tool(): rlm = PredictRLM( "prompt -> answer", sub_lm=mock_lm, - max_iterations=2, + max_iterations=3, sandbox_backend="jspi", ) rlm.generate_action = actions @@ -130,8 +143,8 @@ def test_predict_rlm_jspi_timeout_preserves_state_history_and_predict_tool(): prediction = rlm(prompt="exercise deno timeout recovery") assert prediction.answer == "pre-timeout prediction -> post-timeout prediction / 123" - assert len(prediction.trace.steps) == 2 - timeout_step, final_step = prediction.trace.steps + assert len(prediction.trace.steps) == 3 + _, timeout_step, final_step = prediction.trace.steps assert ( "[Timeout] Iteration execution timed out after 0.2s" in timeout_step.untruncated_output ) @@ -140,11 +153,44 @@ def test_predict_rlm_jspi_timeout_preserves_state_history_and_predict_tool(): assert final_step.output == ( "FINAL: {'answer': 'pre-timeout prediction -> post-timeout prediction / 123'}" ) - second_history = str(actions.calls[1]["repl_history"]) + second_history = str(actions.calls[2]["repl_history"]) assert "[Timeout] Iteration execution timed out after 0.2s" in second_history assert "first predict: pre-timeout prediction" in second_history +@pytest.mark.integration +def test_jspi_timeout_during_async_sleep_preserves_state_and_recovers(monkeypatch): + import predict_rlm.execution_timeout as execution_timeout + from predict_rlm.backends import JspiBackend + from predict_rlm.execution_timeout import RecoverableExecutionTimeout + + monkeypatch.setattr( + execution_timeout, + "DEFAULT_RECOVERABLE_EXECUTION_TIMEOUT_GRACE_SECONDS", + 2, + ) + interpreter = JspiBackend(preinstall_packages=False) + try: + interpreter.execute( + "import asyncio, signal\nsaved = 42\n" + "previous_sigint = signal.getsignal(signal.SIGINT)" + ) + result = interpreter.execute( + "print('before sleep')\nawait asyncio.sleep(0.3)", + timeout=0.05, + ) + assert isinstance(result, RecoverableExecutionTimeout) + assert result.stdout == "before sleep\n" + assert ( + interpreter.execute( + "assert signal.getsignal(signal.SIGINT) is previous_sigint\nprint(saved)" + ) + == "42\n" + ) + finally: + interpreter.shutdown() + + @pytest.mark.asyncio async def test_nonfinite_model_deadline_fails_before_execution(): from dspy.primitives.repl_types import REPLHistory From cd265ffe9c099682a1767ca031ae074868f84e50 Mon Sep 17 00:00:00 2001 From: Emile Riberdy Date: Tue, 8 Sep 2026 17:35:18 +0000 Subject: [PATCH 5/5] docs(predict-rlm): inventory repository test suites --- docs/tests.md | 542 ++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 542 insertions(+) create mode 100644 docs/tests.md diff --git a/docs/tests.md b/docs/tests.md new file mode 100644 index 00000000..9d024080 --- /dev/null +++ b/docs/tests.md @@ -0,0 +1,542 @@ +# Test suites and coverage inventory + +This document maps the repository's tests to the features and failure boundaries +that they protect. It distinguishes package regressions, example-specific tests, +local subprocess contracts, and tests requiring external software or services. + +The inventory describes **implemented test coverage**, not everything an API +might support and not a line- or branch-coverage guarantee. + +## Contents + +- [Discovery and inventory totals](#discovery-and-inventory-totals) +- [Execution tiers and commands](#execution-tiers-and-commands) +- [Shared backend contract matrix](#shared-backend-contract-matrix) +- [Package suite inventory](#package-suite-inventory) +- [Example suite inventory](#example-suite-inventory) +- [Cross-cutting feature coverage](#cross-cutting-feature-coverage) +- [Fixtures, isolation, and external prerequisites](#fixtures-isolation-and-external-prerequisites) +- [CI and interpreting results](#ci-and-interpreting-results) +- [Coverage limits](#coverage-limits) +- [Maintaining the suite and this inventory](#maintaining-the-suite-and-this-inventory) + +## Discovery and inventory totals + +The root [pytest configuration](../pyproject.toml) sets `testpaths = ["tests"]`. +Consequently, running `pytest` or `make test` without explicit paths does **not** +collect the example-local test directories. + +| Collection root | Test modules | Test functions | Collected cases | Purpose | +| --- | ---: | ---: | ---: | --- | +| `tests/` | 45 | 315 | 388 | Default package regression suite: PredictRLM, execution backends, Codex LM, and RLM-GEPA. | +| `examples/terminal_bench/tests/` | 8 | 136 | 136 | Terminal-Bench agents, controller adapters, local runner, scoring, and harness integration seams. | +| `examples/spreadbench/tests/` | 6 | 42 | 43 | Spreadsheet recalculation/rendering and example evaluation configuration. | +| `examples/appworld/tests/` | 1 | 47 | 47 | AppWorld task/session integration, evaluation, scoring, and fixture-backed project behavior. | + +These are inventory snapshots collected with the root project's **all-extras** +environment. Parametrization expands functions into cases. Counts include cases +that subsequently skip because a capability, binary, or service is unavailable. +They are not minimum-count acceptance criteria. Support files such as `conftest.py` +are not included in the test-module count. + +All paths and commands below are relative to the repository root unless stated +otherwise. The per-module tables list **collected cases**, not function counts. + +## Execution tiers and commands + +### Setup and ordinary runs + +The root project requires Python 3.11+ and uses `uv`. The Terminal-Bench example +project declares Python 3.12+. + +```bash +uv sync --all-extras + +# Complete default package suite, including locally enabled integration cases. +uv run --all-extras pytest + +# Compact output without the configured testdox/pass-output verbosity. +uv run --all-extras pytest -o addopts= -q + +# Inventory only: imports and collects tests without executing them. +uv run --all-extras pytest --collect-only -q -o addopts= +``` + +Deno v2 is a default Python dependency, but JSPI execution still requires the +`deno` executable to be available. Initial Pyodide and sandbox-package setup can +require network access. **No LM API credentials are required for the package +regressions**: model responses are scripted or transport endpoints are replaced +with local/fake implementations. + +### Marker model + +Markers are independent axes, not a directory hierarchy: + +| Marker | Meaning | Important distinction | +| --- | --- | --- | +| `integration` | Deno/Pyodide execution or the real Docker Sandboxes service in the package suite. | Local CPython subprocesses are not automatically integration tests. Example suites also use this marker for their own integrations, such as LibreOffice. | +| `sbx` | SBX/supervisor subsystem selection; the `sbx` extra supplies `websockets`. | Most non-integration SBX cases use local supervisors or fakes, not a real sandbox service. | +| `gepa` | RLM-GEPA subsystem selection. | Covers deterministic optimization/evaluation mechanics, not a live optimization experiment. | +| `codex_lm` | Codex LM subsystem selection. | Uses simulated provider streams and local servers, not authenticated Codex calls. | +| `local` | Timing-sensitive or large-payload cases excluded from CI. | Local Make targets include these unless explicitly deselected. | + +Do not infer markers from a filename. For example, `test_sbx_pool.py` contains +both unmarked and SBX-marked tests, while `test_lm_config.py` is GEPA-marked. +The inventory tables show the actual routing observed during collection. + +### Make targets + +The [root Makefile](../Makefile) supplies these selections: + +| Target | Pytest selection | Dependency/runtime scope | +| --- | --- | --- | +| `make test` | No marker filter | Runs with the environment selected by `uv run`; does not request all extras. | +| `make test-unit` | `not integration` | Requests all extras. Includes host logic, local processes, and local WebSocket seams. | +| `make test-integration` | `integration` | Requests all extras; combines JSPI and opt-in real SBX cases. | +| `make test-core` | `not integration and not sbx and not gepa and not codex_lm` | Requests no optional extras; includes local CPython processes. | +| `make test-sbx` | `sbx and not integration` | Requests the `sbx` extra; local supervisor/pool coverage. | +| `make test-gepa` | `gepa and not integration` | Requests the `gepa` extra. | +| `make test-codex-lm` | `codex_lm and not integration` | Requests the `codex-lm` extra. | +| `make test-integration-jspi` | `integration and not sbx` | Real JSPI/Deno package cases. | +| `make test-integration-sbx` | `integration and sbx` | Requests the `sbx` extra and sets `PREDICT_RLM_RUN_SBX_TESTS=1`; requires CLI/login/service access. | +| `make test-local` | `local` | Requests all extras; explicitly selects the local-only cases. | + +With all extras installed, the package cases partition as follows. The CI column +also excludes `local`; prerequisite and capability skips can still reduce the +number actually executed. + +| Selection | Collected locally | Selected by current CI | +| --- | ---: | ---: | +| Core | 167 | 166 | +| SBX, non-integration | 63 | 60 | +| GEPA, non-integration | 43 | 43 | +| Codex LM, non-integration | 46 | 46 | +| JSPI integration | 47 | 46 | +| Real SBX integration | 22 | No standing job | + +**Bootstrap exception:** the five Docker bootstrap cases are currently unmarked. +They appear in core/non-integration collections but skip unless +`PREDICT_RLM_RUN_BOOTSTRAP_DOCKER_TESTS=1`. If that environment variable is set, +`make test-unit` can launch Docker builds; `not integration` is not a universal +"no external binaries" guarantee. + +### Focused runs + +Use pytest directly for paths, node IDs, and expression filters: + +```bash +uv run --all-extras pytest tests/test_predict_rlm.py -q +uv run --all-extras pytest tests/test_small_kernel.py -k cancellation -q +uv run --all-extras pytest tests/test_files.py tests/test_file_sync.py -q +uv run --all-extras pytest tests/runtime_contracts -m 'not integration' -q +uv run --all-extras pytest tests/codex_lm -q +uv run --all-extras pytest tests/test_rlm_gepa.py -k patch -q +``` + +Selecting a path does not remove its integration requirements. For example, the +file transfer tests above execute a real JSPI sandbox. + +## Shared backend contract matrix + +The canonical shared runtime suite is [tests/runtime_contracts](../tests/runtime_contracts/). +Its [backend factories](../tests/runtime_contracts/backends.py) and +[fixture](../tests/runtime_contracts/conftest.py) run the same contracts through +four concrete configurations: + +| Matrix ID | Actual boundary | Selection | Prerequisites | +| --- | --- | --- | --- | +| `jspi` | Deno subprocess and Pyodide/WASM | JSPI integration | Deno v2; sandbox initialization/package availability. | +| `python-runner/direct-process` | Real local CPython supervisor/runner | Core | Local Python; no SBX service or WebSocket transport. | +| `sbx/local-websocket` | Real local supervisor payload over a loopback WebSocket, driven through `SbxBackend` | SBX | `websockets`; no Docker Sandboxes service. | +| `sbx` | Actual Docker Sandboxes environment and WebSocket supervisor | Real SBX integration | `sbx` CLI, login/service access, and explicit opt-in. | + +`RuntimeHandle` normalizes result shapes and lifecycle operations, not behavior. +For JSPI, the matrix's reset operation shuts down the interpreter; the next +execution recreates its state. Other matrix backends use their reset API. +The fixture shuts down each handle after its case. + +### Covered shared features + +| Contract | JSPI | Direct | Local SBX WebSocket | Real SBX | +| --- | --- | --- | --- | --- | +| Execute code, retain state, reset state | Covered | Covered | Covered | Opt-in | +| Serialize concurrent top-level execute requests | Covered | Covered | Covered | Opt-in | +| Normalize Python/REPL/unlabelled code fences | Covered | Covered | Covered | Opt-in | +| Report user/syntax errors and execute again | Covered | Covered | Covered | Opt-in | +| Preserve stdout on a user-code error | Covered | Not asserted: exception has no partial-output contract here | Covered | Opt-in | +| Return structured `SUBMIT` output | Covered | Covered | Covered | Opt-in | +| Defer submission finalization through the interpreter API | Unsupported/skipped | Covered | Unsupported/skipped | Unsupported/skipped | +| Mount, create directories, list, write, and collect files | Covered | Covered | Covered | Opt-in | +| Preserve timeout output and recover for another execution | Covered | Covered | Covered | Opt-in | +| Round-trip list/dict/null/text host-tool results | Covered | Covered | Covered | Opt-in | +| Execute synchronous and asynchronous host tools concurrently | Covered | Unsupported/skipped: callbacks are serial | Covered | Opt-in | +| Recover from host-tool exceptions and tool deadlines | Covered | Covered | Covered | Opt-in | +| Round-trip a roughly 950 KB host-tool request | Local-only | Local-only | Local-only | Local-only and opt-in | + +"Covered" identifies an implemented assertion, not a promise that every +configuration has been exercised on a given machine. Unsupported capabilities +are declared in `RuntimeSpec.unsupported`; they are not successful test coverage. +The Direct-only deferred-submit row concerns the **interpreter API**, not the +higher-level PredictRLM submit-confirmation loop. + +| Module | Cases | Core feature and covered behavior | +| --- | ---: | --- | +| [test_execution_contract.py](../tests/runtime_contracts/test_execution_contract.py) | 44 | Execution/state/reset; concurrent request serialization; code fences; user and syntax error recovery; partial error output where supported; `SUBMIT` and deferred finalization; file roundtrips; recoverable execution deadlines with stdout/stderr. | +| [test_tool_contract.py](../tests/runtime_contracts/test_tool_contract.py) | 24 | Host-tool value shapes; genuine overlap of sync/async callbacks where supported; large-payload transport; recovery after tool errors; bounded timeout/recovery while multiple tools are pending. The deadline reproduction runs in a child process so a stalled backend cannot hang the test indefinitely. | + +## Package suite inventory + +Routing labels below refer to the marker selections above. "Mixed" tests can +combine deterministic fakes with genuine filesystem or subprocess behavior; +it does not mean they contact a live model provider. + +### RLM execution, inputs, callbacks, and skills + +| Module | Cases | Routing | Core feature and covered behavior | +| --- | ---: | --- | --- | +| [test_predict_rlm.py](../tests/test_predict_rlm.py) | 12 | Core + JSPI | Action-loop execution and recovery; nested sandbox-defined Pydantic schemas and an `items` output-name collision; defaulted collection nullability; invalid custom-predictor null output; missing LM errors; concurrent LM-context isolation/restoration; submit confirmation and invalidation by intervening work; sync/async fatal error propagation; no double-charging iteration usage. Real sandbox cases use scripted `DummyLM` responses. | +| [test_empty_code_retry.py](../tests/test_empty_code_retry.py) | 3 | Core | Invalid actions recover through the chat-to-JSON adapter fallback, persistently invalid responses exhaust recovery, and empty custom-predictor code is rejected before execution. | +| [test_in_context.py](../tests/test_in_context.py) | 12 | Core | `CtxStr` adapter precedence and ambiguity; delimiter collisions; prepared versus final bound values; concurrent and sequential run-local predictor/signature isolation; invalid prompt-hook returns; rejection of non-string, optional, and output `CtxStr` declarations. These are host-side prompt construction contracts, not LM comprehension tests. | +| [test_callbacks.py](../tests/test_callbacks.py) | 7 | Core + JSPI | Paired iteration events and shared call IDs; end events on execution failure; awaiting asynchronous handlers; handling async callbacks on the sync path; handler failure isolation; delivery of actual sandbox output to callbacks. | +| [test_rlm_skills.py](../tests/test_rlm_skills.py) | 4 | Core + JSPI | Multiple skills execute together in an RLM run; duplicate tool/module ownership is rejected; a skill cannot silently replace a user tool. Package installation is covered separately below. | + +### Runtime kernel, adapters, files, and workspaces + +| Module | Cases | Routing | Core feature and covered behavior | +| --- | ---: | --- | --- | +| [test_small_kernel.py](../tests/test_small_kernel.py) | 32 | Core + SBX | Invocation and resource ownership: acquisition failure/retry, cancellation and worker draining, adapter/session finalization order, preservation of primary errors, release failure reporting, fixed mount declarations, and SyncedFile temporary-file lifetime. Strict evidence cannot publish success before required finalization or after emit/flush/close failure. Cancellation produces paired terminal evidence. Most cases use purpose-built sessions/backends and real synchronization primitives rather than a sandbox process. | +| [test_adapter_contracts.py](../tests/test_adapter_contracts.py) | 13 | Core | Adapter specificity; duplicate/conflicting host mounts; explicit relative destinations; duplicate destination and input/output overlap rejection; sorted/filtered glob preparation; empty-glob and escaping-symlink rejection; output reservation fallback; exclusion of stale files and out-of-reservation output paths. | +| [test_external_input_adapter_contracts.py](../tests/test_external_input_adapter_contracts.py) | 3 | JSPI + SBX + real SBX | A stateless custom input adapter handles interleaved real JSPI invocations; incompatible service requirements are rejected before reused/pool acquisition; an opt-in owned SBX case enforces a read-only external mount. | +| [test_files.py](../tests/test_files.py) | 2 | JSPI | PredictRLM input/output `File` handling with actual bytes: binary transformation and source preservation; hiding host output destinations from sandbox inputs; directory enumeration via `File.from_dir()` into `list[File]`; discovery of generated output files omitted from the submitted list; exclusion of stale host outputs. A scalar input `File` is not used as a directory mount. | +| [test_file_sync.py](../tests/test_file_sync.py) | 4 | JSPI | Owned PredictRLM SyncedFile operation performs binary writeback visible to later calls and removes temporary files. Backend-local cases cover async tool failure/recovery, no-writeback mode with a retained custom host directory, and missing files. This deliberately distinguishes the portable tool-operation path from backend-local SyncedFile handling. | +| [test_workspace.py](../tests/test_workspace.py) | 11 | Core + JSPI + real SBX | Mirror lifecycle and synchronization through maintained backends; conflict detection before writes; preservation of host changes; oversized/skipped file safety; host/root symlink rejection; manifest failures must not delete host files; one workspace item's conflict must not prevent other items from finalizing. Includes a local supervisor seam and opt-in owned SBX execution. | + +The kernel's safety cases protect several separate transitions: failure while +opening an adapter, acquiring a session, installing/binding runtime inputs, +running generated code, finalizing resources, and releasing the backend. A +successful execution test is not a replacement for these ownership boundaries. + +### Backend-specific execution and transport + +These complement the shared matrix; they are not separate copies of its basic +execute/state/tool/file scenarios. + +| Module | Cases | Routing | Core feature and covered behavior | +| --- | ---: | --- | --- | +| [test_interpreter.py](../tests/test_interpreter.py) | 15 | JSPI | Submit defaults/errors; multi-block fence extraction and inline backticks; nested Pydantic/tool serialization, null values, and string fallback for otherwise unserializable values; nested signature schemas; serialization-error recovery; ignoring late tool responses; reconstructed model extra fields and `items` access; fatal process deadline; workspace flush before cancellation shutdown. Uses real Deno/Pyodide. | +| [test_direct_python_backend.py](../tests/test_direct_python_backend.py) | 3 | Core | Real local runner path virtualization: sandbox paths do not poison subsequent executions; directory contents are copied into the sandbox namespace; a regular `Path` global survives timeout recovery with usable path operations. | +| [test_sbx_interpreter.py](../tests/test_sbx_interpreter.py) | 40 | SBX + real SBX | Native supervisor snapshot/recovery behavior, subprocess stdin isolation, child stdout/stderr attribution, runner exit survival, timeout enforcement despite user exception handlers, tool-reader handoff, late-call quarantine, staging-root ownership, reset, SyncedFile writeback, WebSocket authentication, independent reusable supervisors, Pydantic reconstruction/validation, async interruption, post-hook failure precedence, fixed invocation policy, and cancellation-safe worker cleanup. Most cases use local payload processes; the persistent/re-attach/destroy lifecycle is an opt-in real SBX case. | +| [test_sbx_pool.py](../tests/test_sbx_pool.py) | 20 | Core + SBX | Exclusive leasing/reset; startup and shutdown races; failed reset/replacement/prewarm must not lose capacity or requeue retired interpreters; cancellation-safe retirement and loop migration; waiter cancellation; release after configuration failure; fixed session-policy validation; restart after shutdown. Mixes fake-interpreter synchronization contracts and local supervisor-backed pool cases, not real SBX provisioning. | +| [test_jspi_async_operations.py](../tests/test_jspi_async_operations.py) | 7 | Core | Synthetic async JSPI seams: interrupt until execution quiesces, host-task ownership through recoverable timeout/cancellation, quarantine before the next iteration, and post-hook behavior after fatal/cancelled/failed execution. The filename does not make these Deno integration tests. | +| [test_interpreter_io.py](../tests/test_interpreter_io.py) | 4 | Core | Real OS pipes around a synthetic backend: sync/async UTF-8 backpressure without dropped bytes; partial stdout cannot defeat a request deadline; buffered bytes survive that deadline; stdout EOF must not block draining a still-live stderr pipe. | +| [test_supervisor_client.py](../tests/test_supervisor_client.py) | 2 | Core | Synthetic supervisor frames: discard stale responses/errors until the correct request ID arrives; fail cleanly when the resynchronization limit is exhausted. | +| [test_response_id_resync.py](../tests/test_response_id_resync.py) | 6 | Core | JSPI response multiplexing: stale response rejection in sync/async paths, bounded resynchronization, file-operation response routing, and distinguishing tool calls from stale top-level responses. Uses controlled transport frames, not a Deno process. | +| [test_iteration_execution_timeout.py](../tests/test_iteration_execution_timeout.py) | 3 | Core + JSPI | Bounded failure of silent JSPI timeout recovery; a real PredictRLM deadline preserves state/history/output and permits later `predict()` use; non-finite LM-selected deadlines fail before executing code. | +| [test_tool_call_timeout.py](../tests/test_tool_call_timeout.py) | 3 | Core | Hung asynchronous tools return bounded errors; a timed-out synchronous worker does not poison later executor work; evidence-wrapped sync tool deadlines return without waiting for the still-live worker. Uses real tasks/threads with controlled backend seams. | +| [test_runtime_hooks.py](../tests/test_runtime_hooks.py) | 3 | Core | Real local supervisor hook registration, before/after events for file/subprocess operations, clearing hooks, suppression of internal capture operations, and error events for failed user operations. | +| [test_skill_package_integration.py](../tests/test_skill_package_integration.py) | 2 | JSPI + real SBX | A skill installs `python-slugify`, imports it in generated code, and returns a result through PredictRLM. Covers JSPI and an opt-in real SBX pool; not a fake package-list assertion. | +| [test_bootstrap_controller.py](../tests/test_bootstrap_controller.py) | 5 | Core selection; Docker opt-in | Build fixture images to exercise Python/venv/controller bootstrap on Alpine, Python 3.13 slim, and Ubuntu without Python; reject unsupported BusyBox package management and a non-root environment needing privileged repair. Docker builds install dependencies and can require network access. | + +### Traces, telemetry, and cancellation evidence + +| Module | Cases | Routing | Core feature and covered behavior | +| --- | ---: | --- | --- | +| [test_trace.py](../tests/test_trace.py) | 9 | Core | Replacing an in-progress exported trace; base64 sanitization without mutating the full in-memory trace; proposer projection that keeps behavioral evidence but removes accounting/raw fields; strict-evidence projection; grouping predict calls by signature/instructions/model; nested predict/tool collector isolation; history-delta usage and cache-hit accounting; completion metadata/cache statistics. | +| [test_trace_on_cancellation.py](../tests/test_trace_on_cancellation.py) | 3 | Core | Real traced-loop handling with a controlled interpreter: `KeyboardInterrupt` and `CancelledError` preserve completed and pending iteration evidence; an error while building the trace must not replace the original cancellation. | +| [test_telemetry.py](../tests/test_telemetry.py) | 3 | Core | Persisted JSONL spans, timing and trace/parent correlation; redaction of known secret keys/values in telemetry and debug output; stable candidate hashing across equivalent dictionaries and distinct hashes for different candidates. | +| [test_telemetry_analyzer.py](../tests/test_telemetry_analyzer.py) | 5 | GEPA | Failure precedence across lifecycle, timeout, and model-output evidence; row/event precedence; joining task traces with telemetry artifacts; infrastructure-excluded scores; missing or unrelated evidence remains unknown rather than being confidently misclassified. | + +### Optimization and evaluation + +| Module | Cases | Routing | Core feature and covered behavior | +| --- | ---: | --- | --- | +| [test_rlm_gepa.py](../tests/test_rlm_gepa.py) | 35 | GEPA | Candidate acceptance/significance; group-aware sampling; project seed/component validation; merge eligibility, pair deduplication, and stronger-base selection; capped/balanced disagreement evidence plus shared-success guardrails; reject insufficient evidence or non-improving children; preserve other base components; sidecar resume; correct raw/logical spend; numerical reporting and checkpoint precedence/repair; evaluation timeout/cancellation; failure/structured trace artifacts; unique write-once artifact namespaces on resume; no-op patch preservation; proposer-visible evidence serialization. LM/proposer outputs are controlled rather than generated by a live optimization run. | +| [test_rlm_gepa_patch_merge_costs.py](../tests/test_rlm_gepa_patch_merge_costs.py) | 2 | GEPA | Exact patch text survives artifact persistence, including whitespace/Unicode; main and sub-LM proposer costs both reach reports; legacy patch roles sharing an operation ID must not collapse distinct charges. | +| [test_lm_config.py](../tests/test_lm_config.py) | 1 | GEPA | A constructed GEPA LM retries simulated LiteLLM rate-limit failures and eventually returns success. Environment validation and waiting are replaced; this does not contact a provider. | + +The optimizer suite checks the mechanics of evaluation and instruction evolution. +It does not establish that an optimized instruction improves a real benchmark or +that a particular model follows the proposer prompt correctly. + +### Codex LM transport, authentication, and usage + +[tests/codex_lm/conftest.py](../tests/codex_lm/conftest.py) marks that directory as +`codex_lm`. It clears/disables disk caching for isolation, disables production +retry delays by default, isolates the auth home, and supplies controlled stream +events and LM instances. Retry tests opt back into their required retry behavior. +The root WebSocket module is marked separately. + +| Module | Cases | Routing | Core feature and covered behavior | +| --- | ---: | --- | --- | +| [test_auth.py](../tests/codex_lm/test_auth.py) | 12 | Codex | Explicit legacy-auth opt-in; private credential persistence/removal; enable/disable state; profile-name validation and slug collisions; explicit/environment/active-profile precedence; rotation skips disabled accounts; long-lived LM refresh after account-state changes; redacted status metadata. Uses temporary profiles, not real accounts. | +| [test_auto_retry.py](../tests/codex_lm/test_auto_retry.py) | 2 | Codex | Retry a stalled stream using another enabled rotation profile; expose `CodexStreamError` after retry exhaustion. | +| [test_build_request.py](../tests/codex_lm/test_build_request.py) | 2 | Codex | Request conversion does not mutate reusable constructor reasoning configuration; request-scoped proxy environment changes are restored. | +| [test_cli.py](../tests/codex_lm/test_cli.py) | 8 | Codex | Interception of supported DSPy OpenAI construction; rejection of unsupported models; other providers pass through; child argv/exit status behavior; disabled-profile usage avoids live fetches; login uses an isolated Codex home and preserves the login exit code. CLI invocations are exercised with controlled credentials/login seams, not an actual login. | +| [test_concurrent_cache.py](../tests/codex_lm/test_concurrent_cache.py) | 1 | Codex | Concurrent requests keep their results and cache keys separate; repeated requests hit the correct cache entry instead of transport. Does not promise defensive copying of caller-mutated cached objects. | +| [test_forward.py](../tests/codex_lm/test_forward.py) | 3 | Codex | HTTP stream assembly; cached-input/output pricing; cached calls do not double-charge usage or alter fresh history accounting; asynchronous DSPy prediction exposes parsed answers and billable usage. | +| [test_stream_errors.py](../tests/codex_lm/test_stream_errors.py) | 4 | Codex | Failed, incomplete, explicit-error, and truncated streams surface errors rather than being returned as successful responses. | +| [test_stream_heartbeat.py](../tests/codex_lm/test_stream_heartbeat.py) | 3 | Codex | A silent async stream times out; async HTTP completion does not wait for connection closure; sync HTTP consumption stops at the completion event. Separate sync/async paths both have a terminal-boundary assertion. | +| [test_stream_redaction.py](../tests/codex_lm/test_stream_redaction.py) | 1 | Codex | Stream failure diagnostics do not expose the supplied credential values. | +| [test_usage.py](../tests/codex_lm/test_usage.py) | 3 | Codex | Derive remaining credit and nested model limits; keep live usage windows distinct from model-specific limits; preserve profile display names without leaking raw secret payloads. | +| [test_ws_lm.py](../tests/codex_lm/test_ws_lm.py) | 2 | Codex | Retry-scoped turn state does not leak to the next invocation; exhausted WebSocket transport falls back to HTTP and remains on HTTP for subsequent calls. Uses controlled transport seams, including a local HTTP response path. | +| [test_codex_ws_lm.py](../tests/test_codex_ws_lm.py) | 5 | Codex | Real loopback WebSocket server: prewarm before generation, preserve server error details, isolate turn state per forward, finish synchronously while the server keeps the connection open, and classify HTTP 401 as expired authentication. This is local protocol validation, not the provider service. | + +## Example suite inventory + +Example suites are collected only when selected explicitly. They retain some +source-text, prompt-wording, defaults, and wiring checks that are intentionally +absent from the reduced package suite. Inventorying them does not mean those +checks meet the package suite's preferred behavioral-test standard. + +The following root-environment collection command was checked without running +example tests: + +```bash +uv run --all-extras pytest --collect-only -q -o addopts= \ + examples/terminal_bench/tests \ + examples/spreadbench/tests \ + examples/appworld/tests +``` + +To execute a suite, remove `--collect-only` or select its path separately. Passing +collection is not proof that its execution-time binaries or optional packages are +available, nor that the example suite currently passes. + +### Terminal-Bench: 136 cases in eight modules + +```bash +uv run --all-extras pytest examples/terminal_bench/tests -q +``` + +The [example Makefile](../examples/terminal_bench/Makefile) also provides `make test` +when invoked from that example directory. Its `make setup` provisions a separate +Terminal-Bench environment, and `make smoke` runs a synthetic three-task scoring +scenario. These are not substitutes for a real benchmark run. See the +[setup script](../examples/terminal_bench/scripts/setup_terminal_bench.sh) for harness environment installation. + +| Module | Cases | Core feature and covered behavior | +| --- | ---: | --- | +| [test_container_runner.py](../examples/terminal_bench/tests/test_container_runner.py) | 19 | Python supervisor request/reset/shutdown mapping; process-clear/shutdown races; code-fence and execution-timeout forwarding; real local runner stdout, timeout recovery, child exit, and stdin isolation; restart diagnostics; structured error mapping; backend-routed file operations; bounded host-tool timeout. Despite the filename, these use fakes and local Python subprocesses, not a container service. | +| [test_gepa_project.py](../examples/terminal_bench/tests/test_gepa_project.py) | 59 | Project/config/CLI contracts; task timeouts/resources; Harbor controller locality selection; supplied Daytona/controller requirements; local shell file transfer and simulated SSH/SBX/SDK adapters; source-bundle/bootstrap layout; remote-root preservation; tracked-file packaging and tar traversal rejection; selective transient setup retries; phase-event aggregation; official task-cache timeout lookup and bounded failure diagnostics; Harbor result/CTRF parsing; GEPA scoring, traces, and LM/agent argument propagation. Most remote/harness operations are fake; local filesystem, archive, and shell paths are exercised. | +| [test_harbor_agent.py](../examples/terminal_bench/tests/test_harbor_agent.py) | 20 | Remote agent payload/environment boundaries; answer sentinel parsing; status/trace persistence and download on success/cancellation; debug streaming and shutdown failure handling; confirmation callback reconstruction; bootstrap and opaque auth upload; Harbor context metadata; setup/agent phase events. Uses simulated remote environments. One context-model case skips if its optional Harbor import is unavailable. | +| [test_runner.py](../examples/terminal_bench/tests/test_runner.py) | 15 | Real local JSON-RPC runner: state/reset, predict/tool and image data URL roundtrips, path behavior, host errors, child-output attribution, runner exit recovery, non-swallowable/native-blocking deadlines, and termination of generated child processes. Also contains a source-text payload-sharing check. | +| [test_scoring.py](../examples/terminal_bench/tests/test_scoring.py) | 8 | Full/partial rewards and CTRF detail precedence; evaluator exception/timeout classification; verified pass evidence overrides timeout placeholders where appropriate; numeric GEPA objective scores. | +| [test_setup_wiring.py](../examples/terminal_bench/tests/test_setup_wiring.py) | 4 | Static package dependencies/entry points, setup-script content, and Make targets. These do not run installation or prove a provisioned environment works. | +| [test_smoke.py](../examples/terminal_bench/tests/test_smoke.py) | 1 | Runs a local synthetic scoring script and distinguishes all-pass, partial, and all-fail outcomes. Does not execute Terminal-Bench tasks against an LM. | +| [test_tbench_agent.py](../examples/terminal_bench/tests/test_tbench_agent.py) | 10 | Agent names/factory, mocked PredictRLM construction, custom signature and confirmation configuration, Codex installation order/error hints, trace export, and rejection of wrapper tools. These are adapter/wiring contracts, not live agent performance tests. | + +### Spreadbench: 43 cases in six modules + +```bash +uv run --all-extras pytest examples/spreadbench/tests -q +``` + +The root `examples` extra supplies workbook/PDF-related Python dependencies, +including `openpyxl`, `formulas`, and PyMuPDF. Real rendering additionally requires +**LibreOffice/`soffice` and Poppler's `pdftoppm` on PATH**. The render module skips +if either binary is absent. Recalculation has LibreOffice-dependent cases, +including an `integration`-marked rescue path; absence of that marker on another +case does not imply no external binaries. + +| Module | Cases | Core feature and covered behavior | +| --- | ---: | --- | +| [test_eval_sbx_pool.py](../examples/spreadbench/tests/test_eval_sbx_pool.py) | 7 | Backend/pool/logging CLI configuration, invalid pool/backend combinations, and fake pool/PredictRLM argument and lifecycle wiring. Does not provision SBX or execute JSPI. | +| [test_gepa_telemetry.py](../examples/spreadbench/tests/test_gepa_telemetry.py) | 2 | Case start/end telemetry and preservation of a host-tool error span when best-effort recalculation catches an exception. Model execution and scoring are controlled. | +| [test_instruction_prompt.py](../examples/spreadbench/tests/test_instruction_prompt.py) | 1 | Static instruction framing: natural-language requests describe workbook edits, preserve existing values, and do not solicit a prose answer. Not an instruction-following evaluation. | +| [test_recalculate.py](../examples/spreadbench/tests/test_recalculate.py) | 15 | Actual generated workbooks and formula caches; target discovery/resolution counts; formula-library evaluation; preservation of non-formula content; additive/no-op behavior; missing/no-formula inputs; winner/tie precedence; missing-library/error/timeout fallback; LibreOffice rescue of incomplete results. Some failure/fallback seams are monkeypatched. | +| [test_recalculate_hang.py](../examples/spreadbench/tests/test_recalculate_hang.py) | 2 | A checked-in full-column-reference workbook is recalculated in a bounded child process to defend against hangs; no-formula workbook calls emit host-tool telemetry. The hang reproduction disables LibreOffice and bounds formula-worker termination. | +| [test_render.py](../examples/spreadbench/tests/test_render.py) | 16 | Real workbook-to-PNG rendering and data URIs; string/path inputs; cell ranges and sheet-qualified selection; missing files/bad sheets/empty ranges; wrapper error conversion and registration/forwarding checks. Entire module is gated on LibreOffice and Poppler. | + +The [Spreadbench README](../examples/spreadbench/README.md) also describes live +evaluation prerequisites. Dataset downloads, provider keys, and live LM runs are +not implied by these example unit/tool tests. + +### AppWorld: 47 cases in one module + +```bash +uv run --all-extras pytest examples/appworld/tests -q +``` + +[conftest.py](../examples/appworld/tests/conftest.py) adds the example package to +`sys.path`. Tests use tiny fixture datasets, generated task assets, fake workers, +and controlled RLM results; they do not launch the real AppWorld runtime. + +| Module | Cases | Core feature and covered behavior | +| --- | ---: | --- | +| [test_appworld_smoke.py](../examples/appworld/tests/test_appworld_smoke.py) | 47 | Service/project construction; official ICL manifest loading and runtime demo adaptation; dataset/spec loading and deterministic group-disjoint splits; evaluator/count-derived scoring; worker path/JSON conversion; isolated runtime discovery; session JSON argument validation and EOF/stderr deadlock prevention; hiding model-facing completion APIs and internal fields; persistence before evaluation; completion from several answer shapes without legacy fallback/double completion; task-bound host tools; harness-side scoring; evaluation artifacts/LM construction; Codex CLI setup and error hints. Also includes static prompt/default/wiring assertions. | + +Live AppWorld execution uses a separate `.appworld-venv` because its dependency +stack includes Pydantic v1. The normal PredictRLM environment remains separate. +Follow the [AppWorld README](../examples/appworld/README.md) for runtime/data +setup; those installations and provider credentials are not needed for the +fixture-backed suite described here. + +## Cross-cutting feature coverage + +Use this map when deciding where a new regression belongs. The feature's owner +is more useful than the particular bug report or backend that exposed it. + +| Feature or invariant | Primary coverage locations | +| --- | --- | +| Generated code produces the final typed result | `test_predict_rlm.py`; shared execution contracts; `test_files.py`. | +| Invalid model actions do not reach code execution | `test_empty_code_retry.py`; `test_iteration_execution_timeout.py`; predict-output validation in `test_predict_rlm.py`. | +| Invocation-local prompts, LM contexts, and collector state | `test_in_context.py`; `test_predict_rlm.py`; `test_trace.py`; external adapter interleaving. | +| Adapter specificity, declared path ownership, and output reservations | `test_adapter_contracts.py`; `test_small_kernel.py`; `test_external_input_adapter_contracts.py`. | +| Host changes and files survive failure without silent clobbering/deletion | `test_workspace.py`; `test_files.py`; `test_file_sync.py`; workspace cancellation in `test_interpreter.py`. | +| Cancellation does not release a lease or delete staging while work is live | `test_small_kernel.py`; `test_jspi_async_operations.py`; `test_sbx_pool.py`; `test_sbx_interpreter.py`. | +| Primary exceptions survive failures in finalization, callbacks, or trace building | `test_small_kernel.py`; `test_callbacks.py`; `test_trace_on_cancellation.py`; backend post-hook cases. | +| Stale responses and partial pipe data cannot corrupt/hang the protocol | `test_response_id_resync.py`; `test_supervisor_client.py`; `test_interpreter_io.py`; native supervisor handoff cases. | +| Timeouts are bounded and recoverable where specified | Shared execution/tool contracts; iteration/tool timeout modules; backend-specific native recovery cases. | +| Credentials and raw image/accounting payloads stay out of restricted outputs | Codex auth/redaction/usage tests; `test_telemetry.py`; `test_trace.py`; proposer artifact tests. | +| Cached or resumed work does not inflate or erase spend | Codex forward/cache tests; `test_trace.py`; GEPA cost/reporting and patch-merge cost tests. | +| Instruction patches require evidence and preserve unrelated solved behavior | GEPA acceptance, balanced/shared-success evidence, improvement gate, and base-component preservation tests. | +| Benchmark scores are based on evaluator evidence, not model claims | Example Terminal-Bench scoring and project tests; AppWorld harness-side evaluation tests. | + +## Fixtures, isolation, and external prerequisites + +### Shared support and lifecycle ownership + +- The shared backend fixture always shuts down the runtime it creates. Its + factories separate local Direct, local WebSocket, and real SBX environments. +- Kernel tests use explicit fake sessions/backends and synchronization events to + make cleanup order, acquisition failures, and live-worker ownership observable. + A fake session proving ordering is not a real filesystem-transfer test. +- Files/workspaces use temporary host directories; the owned SyncedFile test + additionally verifies staging cleanup after actual roundtrips. +- Codex fixtures isolate cache/auth state. CLI fixtures restore DSPy LM symbols, + `sys.argv`, and `sys.path` after script execution. HTTP/WebSocket tests use + loopback servers or simulated event streams. +- Example AppWorld uses checked-in tiny datasets and generated task assets. + Spreadbench creates workbooks in temporary directories and retains a captured + huge-range workbook for the bounded-hang reproduction. +- The bootstrap image fixtures live under + [tests/fixtures/bootstrap_controller](../tests/fixtures/bootstrap_controller/). + They are Docker build scenarios, not Python test modules. + +### Explicit external runs + +Real SBX tests create and remove sandbox resources. The persistent lifecycle case +also reattaches to a named sandbox before destroying it. Use an appropriate test +account/environment, not a sandbox containing unrelated work. + +```bash +# Requires an installed sbx CLI and a valid login/service connection. +make test-integration-sbx + +# Requires Docker CLI and a reachable daemon; builds five fixture images. +PREDICT_RLM_RUN_BOOTSTRAP_DOCKER_TESTS=1 \ + uv run --all-extras pytest tests/test_bootstrap_controller.py -q +``` + +Local WebSocket tests need loopback networking but no remote sandbox account. +Deno/Pyodide initialization, skill package installation, and bootstrap Docker +builds can download runtime/package artifacts even though no LM is contacted. + +## CI and interpreting results + +The [test workflow](../.github/workflows/tests.yml) runs: + +- Core tests on Ubuntu with Python 3.11, 3.12, and 3.13. +- SBX, GEPA, Codex LM, and JSPI selections in separate Python 3.12 jobs. +- `and not local` in every pytest job. +- Ruff over `src/` and `tests/`. + +The workflow does not explicitly run the example-local suites, real SBX, or +opt-in bootstrap Docker scenarios. The six collected `local` cases comprise the +four backend variants of the large host-tool payload contract and two native +SBX/supervisor timing/interrupt cases. One of those four payload variants also +requires real SBX opt-in. + +Current CI pytest invocations request `--cov=predict_rlm` and upload coverage +artifacts. Test counts and those artifacts should not be read as a separate +coverage guarantee for `rlm_gepa`, `dspy_codex_lm`, or example packages. + +A recent full package run in the all-extras environment completed with +**358 passed, 30 skipped**. The skips were 21 real-SBX prerequisite cases, five +bootstrap Docker cases, and four unsupported shared-matrix capability cases. +That run also reported DSPy deprecation warnings and two async-stream `aclose` +cleanup warnings. They were not suppressed. This historical execution result is +not an execution claim for the example suites, which were collected separately +for this inventory. + +When reading a result: + +1. Check the selected paths and markers. A green `test-core` run is not a full + package or example run. +2. Use `-rs` to inspect skip reasons. A capability skip is different from a + missing binary, missing optional dependency, or disabled real-service test. +3. Distinguish controlled model responses from live provider behavior and local + SBX supervisor seams from real SBX provisioning. +4. Treat unexpected warnings and incomplete evidence as diagnostics, not as + functionality proven merely because pytest exited successfully. + +## Coverage limits + +- These tests do not measure real model quality, benchmark accuracy, or live + provider authentication/API compatibility. Scripted LMs and local protocol + servers make the relevant software contracts deterministic. +- The backend matrix does not assert uniform capabilities. Direct serial + callbacks, partial-error-output differences, and deferred-submit support are + explicit limitations of the tested interfaces. +- Nested model reconstruction is exercised for sandbox-global model definitions; + this is not a guarantee of reconstruction for arbitrary function-local classes. +- The successful portable SyncedFile roundtrip uses an **owned** PredictRLM JSPI + session. It is not proof that every injected legacy interpreter has the same + reentrant artifact-transfer behavior. +- Workspace/path and credential-redaction tests cover specified safety + boundaries; they are not a comprehensive sandbox security audit or proof that + every possible secret representation is redacted. +- Timeout and overlap assertions defend boundedness/ownership, not throughput + targets. This suite is not a performance benchmark. +- Example source/default/prompt checks prove strings or configuration shape, not + installation success, semantic instruction-following, or a live benchmark run. +- External service/image coverage is absent from a run unless its prerequisites + are enabled. Tests named for a platform-specific path pattern do not establish + native execution on that platform; current CI runs on Ubuntu. + +## Maintaining the suite and this inventory + +### Where to add a test + +1. Find the owning feature in the cross-cutting map and existing module inventory. +2. For shared execution behavior, extend `tests/runtime_contracts/` rather than + copying the scenario into each backend module. +3. Keep backend-specific process, protocol, cancellation, or policy failures in + their owning backend module. Declare actual capability differences explicitly. +4. Keep example harness/config/scoring behavior under the relevant example. +5. Add a permanent case only when a plausible regression would violate an + observable contract. Avoid export/default/docstring/source-text assertions, + mock echoes, and multiple variants exercising the same branch. +6. Preserve genuinely distinct failure transitions: cleanup, primary-error + precedence, data loss, account isolation, and terminal stream completion are + not interchangeable happy-path checks. +7. Bound waits and make teardown own every process, server, thread, lease, or + temporary artifact created by the scenario. Do not inherit a collected test + class just to reuse its fixtures and accidentally execute all its tests again. + +### Updating and checking this document + +When adding, removing, moving, or remarking tests, update the relevant table row, +feature map, prerequisites, and collection snapshot. Do not add tests that pin +this Markdown or require its counts to remain constant. + +```bash +# Recompute the default suite's node IDs and count. +uv run --all-extras pytest --collect-only -q -o addopts= + +# Inspect an exact CI-like selection without running it. +uv run --all-extras pytest --collect-only -q -o addopts= \ + -m 'integration and not sbx and not local' + +# Include skip reasons when executing a selected suite. +uv run --all-extras pytest tests/runtime_contracts -q -rs + +# Run the repository's Python lint check. +uv run ruff check src/ tests/ +``` + +For long-running local benchmark/evaluation work beyond these regression tests, +follow [the local-run runbook](runbooks/long-running-local-runs.md). For execution +ownership and backend boundaries, consult [the architecture](../ARCHITECTURE.md); +for custom inputs and sessions, see [custom adapters](custom-adapters.md) and +[custom path inputs](custom-path-inputs.md).