From c4360bbb59caefce245104b2e1424fab1236af36 Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 20 Oct 2026 20:06:11 +0000 Subject: [PATCH] fix(harness): set PWD to the working directory of local sandbox commands (#47530) * fix(harness): set PWD to the working directory of local sandbox commands LocalSandbox started CLIs in the sandbox workdir but passed the parent's PWD through. OpenCode reads PWD as its project directory, so a litellm.agent() turn worked in the Python process cwd instead of the sandbox. * fix(harness): re-home resumed OpenCode sessions into the sandbox folder OpenCode files each session under the folder it ran in and only streams a resumed session's events from that folder, so a session saved before PWD pointed at the sandbox, or resumed in another sandbox, hung. Resume now runs opencode export and import in the sandbox first. Also types the CLI handler's config so its member accesses are known. * fix(harness): fail an OpenCode resume that cannot move its session, with a timeout The export and import step now has a 30 s bound, and a failure or timeout raises HarnessError instead of leaving the resumed turn waiting forever. A restart of a session the handler already ran skips the step. * fix(harness): resume an OpenCode session elsewhere as a copy and fail the turn when it cannot Importing the exported session moved it, so the original stopped working where it was saved. Resume now imports a copy under fresh ids, filed under the sandbox folder, and continues the copy. A failed or timed-out export or import ends the next turn with runtime_error, as base does for a missing session, instead of raising from agent_resume. * fix(harness): retry a failed OpenCode resume on the next turn and keep resumed copies stable Resume now only records the saved id, and the first turn runs the OpenCode export and import under the caller's timeout. A failed export fails that turn and the next turn tries again. The local sandbox kills a command when its run is cancelled, and the copy names the sandbox folder so OpenCode 1.15.5 and older keep it there. * test(harness): keep subprocess tests out of tests/unit The PWD regression test moves to tests/integration/sdk with an isolated child interpreter, and the cancel test checks collect_output with an in-process process instead of a real one. collect_output takes the sandbox Process protocol. * fix(harness): rewrite only id fields when copying an OpenCode session for resume --- litellm/constants.py | 1 + litellm/harness/handlers/cli_handler.py | 86 +++++++- litellm/harness/sandbox/local.py | 16 +- .../llms/base_llm/harness/transformation.py | 22 ++ .../llms/opencode/harness/transformation.py | 115 +++++++++- .../code_coverage_tests/recursive_detector.py | 1 + tests/harness_e2e/test_harness_e2e.py | 65 +++++- .../integration/sdk/test_local_sandbox_sdk.py | 39 ++++ tests/unit/harness/sandbox/test_local.py | 29 ++- .../opencode/harness/test_transformation.py | 204 +++++++++++++++++- 10 files changed, 554 insertions(+), 24 deletions(-) create mode 100644 tests/integration/sdk/test_local_sandbox_sdk.py diff --git a/litellm/constants.py b/litellm/constants.py index a0dcf44d121..cb50d74f0cf 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -2320,6 +2320,7 @@ HARNESS_STDERR_TAIL_LINES: Final = 40 HARNESS_STREAM_READ_CHUNK_BYTES: Final = 64 * 1024 HARNESS_EVENT_QUEUE_MAX_SIZE: Final = 1024 HARNESS_PROCESS_KILL_GRACE_SECONDS: Final = 5.0 +HARNESS_RESUME_COMMAND_TIMEOUT_SECONDS: Final = 30.0 HARNESS_SNAPSHOT_SKIP_DIRS: Final = frozenset( { ".git", diff --git a/litellm/harness/handlers/cli_handler.py b/litellm/harness/handlers/cli_handler.py index d217d1e5bbb..6320a877e6b 100644 --- a/litellm/harness/handlers/cli_handler.py +++ b/litellm/harness/handlers/cli_handler.py @@ -12,17 +12,27 @@ from __future__ import annotations import asyncio import os from collections import deque -from collections.abc import AsyncIterator, Sequence +from collections.abc import AsyncIterator, Mapping, Sequence +from dataclasses import dataclass from typing import Final from litellm._logging import verbose_logger -from litellm.constants import HARNESS_STDERR_TAIL_LINES, HARNESS_STREAM_READ_CHUNK_BYTES +from litellm.constants import ( + HARNESS_RESUME_COMMAND_TIMEOUT_SECONDS, + HARNESS_STDERR_TAIL_LINES, + HARNESS_STREAM_READ_CHUNK_BYTES, +) from litellm.harness.context import SessionContext -from litellm.harness.errors import HarnessInstallFailed, SandboxError +from litellm.harness.errors import HarnessError, HarnessInstallFailed, SandboxError from litellm.harness.handlers.base import BaseHarnessHandler -from litellm.harness.sandbox.base import Process, Sandbox +from litellm.harness.sandbox.base import CompletedRun, Process, Sandbox from litellm.harness.types import Event -from litellm.llms.base_llm.harness.transformation import BaseCLIHarnessConfig, HarnessSessionSetup +from litellm.llms.base_llm.harness.transformation import ( + BaseCLIHarnessConfig, + HarnessSessionCopy, + HarnessSessionSetup, + HarnessTurnError, +) from litellm.llms.base_llm.harness.utils import decode_json_line, read_skill_files # Link / to a LiteLLM-owned cache dir so a later session can resume. @@ -81,6 +91,52 @@ async def persist_dir(sandbox: Sandbox, link_path: str, cache_subpath: str) -> N ) +@dataclass(frozen=True, slots=True) +class ResumeFailed: + reason: str + + +async def bounded_resume_run( + sandbox: Sandbox, cmd: Sequence[str], env: Mapping[str, str] +) -> CompletedRun | ResumeFailed: + try: + run: Final = await sandbox.run(cmd, env=env, timeout=HARNESS_RESUME_COMMAND_TIMEOUT_SECONDS) + except SandboxError as exc: + return ResumeFailed(str(exc)) + return run if run.exit_code == 0 else ResumeFailed(run.stderr.strip() or f"exit code {run.exit_code}") + + +async def import_session_copy(sandbox: Sandbox, copy: HarnessSessionCopy, env: Mapping[str, str]) -> str | ResumeFailed: + try: + await sandbox.write(copy.path, copy.data) + except SandboxError as exc: + return ResumeFailed(str(exc)) + imported: Final = await bounded_resume_run(sandbox, copy.import_command, env) + return imported if isinstance(imported, ResumeFailed) else copy.session_id + + +async def resumable_session_id( + config: BaseCLIHarnessConfig[object, object], + sandbox: Sandbox, + setup_env: Mapping[str, str], + private_dir: str, + native_session_id: str, +) -> str | ResumeFailed: + export_cmd: Final = config.transform_resume_export_command(native_session_id) + if export_cmd is None: + return native_session_id + exported: Final = await bounded_resume_run(sandbox, export_cmd, setup_env) + if isinstance(exported, ResumeFailed): + return exported + try: + copy: Final = config.transform_resume_copy(exported.stdout, private_dir, sandbox.workdir) + except HarnessError as exc: + return ResumeFailed(str(exc)) + if copy is None: + return native_session_id + return await import_session_copy(sandbox, copy, setup_env) + + async def copy_skills(sandbox: Sandbox, skills: Sequence[str], skills_root: str) -> None: for skill in skills: name = os.path.basename(os.path.realpath(os.fspath(skill))) @@ -89,13 +145,14 @@ async def copy_skills(sandbox: Sandbox, skills: Sequence[str], skills_root: str) class CLIHarnessHandler(BaseHarnessHandler): - config: BaseCLIHarnessConfig + config: BaseCLIHarnessConfig[object, object] def __init__(self, config: BaseCLIHarnessConfig) -> None: super().__init__(config) self._private_dir: str | None = None self._setup: HarnessSessionSetup | None = None self._native_id: str | None = None + self._pending_resume: str | None = None self._proc: Process | None = None async def start(self, ctx: SessionContext) -> None: @@ -119,6 +176,9 @@ class CLIHarnessHandler(BaseHarnessHandler): async def turn(self, ctx: SessionContext, prompt: str) -> AsyncIterator[Event]: if self._setup is None or self._private_dir is None: raise RuntimeError("CLIHarnessHandler.turn() called before start()") + if self._pending_resume is not None: + self._native_id = await self._prepared_session_id(ctx, self._setup, self._private_dir, self._pending_resume) + self._pending_resume = None request = self.config.transform_turn_request(ctx, self._setup, self._private_dir, prompt, self._native_id) argv: Final = list(request.argv) # mutable-ok: Sandbox.exec takes list[str] proc = await ctx.sandbox.exec(argv, env=request.env, cwd=request.cwd) @@ -158,4 +218,16 @@ class CLIHarnessHandler(BaseHarnessHandler): return self._native_id async def resume(self, ctx: SessionContext, native_session_id: str) -> None: - self._native_id = native_session_id + if native_session_id != self._native_id: + self._native_id = native_session_id + self._pending_resume = native_session_id + + async def _prepared_session_id( + self, ctx: SessionContext, setup: HarnessSessionSetup, private_dir: str, saved_id: str + ) -> str: + target: Final = await resumable_session_id(self.config, ctx.sandbox, setup.env, private_dir, saved_id) + match target: + case ResumeFailed(reason=reason): + raise HarnessTurnError(f"could not prepare session {saved_id} for resume: {reason}") + case session_id: + return session_id diff --git a/litellm/harness/sandbox/local.py b/litellm/harness/sandbox/local.py index 10c7303ef27..682e80b33a8 100644 --- a/litellm/harness/sandbox/local.py +++ b/litellm/harness/sandbox/local.py @@ -14,7 +14,7 @@ from typing import Final from litellm.constants import HARNESS_PROCESS_KILL_GRACE_SECONDS from litellm.harness.errors import SandboxError -from litellm.harness.sandbox.base import CompletedRun +from litellm.harness.sandbox.base import CompletedRun, Process from litellm.harness.sandbox.snapshot import snapshot_local _SECRET_PREFIXES: Final = ( @@ -103,7 +103,7 @@ class SubprocessHandle: await self._proc.wait() -async def _read_all(handle: SubprocessHandle) -> tuple[bytes, bytes, int]: +async def _read_all(handle: Process) -> tuple[bytes, bytes, int]: if handle.stdin is not None: handle.stdin.close() stdout, stderr = await asyncio.gather(handle.stdout.read(), handle.stderr.read()) @@ -111,13 +111,16 @@ async def _read_all(handle: SubprocessHandle) -> tuple[bytes, bytes, int]: return stdout, stderr, exit_code -async def collect_output(handle: SubprocessHandle, cmd: Sequence[str], timeout: float | None) -> CompletedRun: - """Close stdin, read stdout/stderr to EOF; kill and raise SandboxError on timeout.""" +async def collect_output(handle: Process, cmd: Sequence[str], timeout: float | None) -> CompletedRun: + """Close stdin, read stdout/stderr to EOF; kill on cancel, and kill and raise SandboxError on timeout.""" try: stdout, stderr, code = await asyncio.wait_for(_read_all(handle), timeout) except asyncio.TimeoutError: await handle.kill() raise SandboxError(f"command timed out after {timeout}s: {cmd[0]}") + except asyncio.CancelledError: + await handle.kill() + raise return CompletedRun( stdout=stdout.decode("utf-8", errors="replace"), stderr=stderr.decode("utf-8", errors="replace"), @@ -180,14 +183,15 @@ class LocalSandbox: self._check_open() if not cmd: raise SandboxError("exec() needs a non-empty command") + work_dir: Final = self._resolve_cwd(cwd) try: proc = await asyncio.create_subprocess_exec( *cmd, stdin=asyncio.subprocess.PIPE, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE, - cwd=self._resolve_cwd(cwd), - env=self.child_env(env), + cwd=work_dir, + env=self.child_env(MappingProxyType({**(env or {}), "PWD": work_dir})), start_new_session=True, ) except (FileNotFoundError, PermissionError) as exc: diff --git a/litellm/llms/base_llm/harness/transformation.py b/litellm/llms/base_llm/harness/transformation.py index d6db3d2ecb1..7eaf6d5aca0 100644 --- a/litellm/llms/base_llm/harness/transformation.py +++ b/litellm/llms/base_llm/harness/transformation.py @@ -68,6 +68,16 @@ class HarnessTurnRequest: cwd: str | None = None +@dataclass(frozen=True) +class HarnessSessionCopy: + """A saved session under fresh ids: written to `path` in the sandbox, then loaded by `import_command`.""" + + session_id: str + path: str + data: bytes + import_command: Sequence[str] + + @dataclass(frozen=True) class HarnessTurnResponse: """What the runtime produced for one turn, after the process exited.""" @@ -128,6 +138,18 @@ class BaseCLIHarnessConfig(BaseHarnessConfig[OptionsT], Generic[OptionsT, Stream ) -> HarnessTurnRequest: """argv / env / stdin for one turn. native_session_id is set after the first turn.""" + def transform_resume_export_command(self, native_session_id: str) -> Sequence[str] | None: + """Command that prints a saved session before resuming one this handler did not run. + + None resumes the saved session as is. A non-zero exit or a timeout fails the turn, and the next turn + tries again. + """ + return None + + def transform_resume_copy(self, exported: str, private_dir: str, workdir: str) -> HarnessSessionCopy | None: + """The copy to import and continue, from the export command's stdout, or None to continue the original.""" + return None + @abstractmethod def create_stream_state(self) -> StreamStateT: """Fresh per-turn parser state.""" diff --git a/litellm/llms/opencode/harness/transformation.py b/litellm/llms/opencode/harness/transformation.py index 144687233e8..46e8b750ca5 100644 --- a/litellm/llms/opencode/harness/transformation.py +++ b/litellm/llms/opencode/harness/transformation.py @@ -8,17 +8,31 @@ repo's own opencode.json cannot redirect model calls. The token is never in argv the config references it with `{file:/token}`. XDG dirs point at a persisted LiteLLM-owned root so the user's opencode config and auth are never read, and the session DB outlives a session for resume. Verified against opencode 1.14.41. + +opencode files each session under the folder it ran in and only streams a resumed session's +events from that folder, so resuming one filed elsewhere (by an older LiteLLM that left PWD +unset, or from another sandbox) hangs. The first turn after a resume runs `opencode export` in +the sandbox, and when the session is filed elsewhere it imports a copy under fresh ids, filed +under the sandbox folder, and continues the copy. Import overwrites the folder of a session id it +already has, so reusing the ids would break the original wherever it was saved. Before 1.15.6, +import files a session under the folder named in its JSON, so the copy names the sandbox folder. +A failed or timed-out export or import fails that turn instead of hanging, and the next turn +tries again. Verified against opencode 1.14.41, 1.15.5, 1.15.6 and 1.19.3. """ from __future__ import annotations import itertools import json -from collections.abc import Mapping, Sequence +import secrets +import string +from collections.abc import Iterator, Mapping, Sequence from dataclasses import dataclass from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final +from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter, ValidationError + from litellm.harness.errors import CapabilityUnsupported, HarnessError, OptionsMismatch from litellm.harness.options import OpenCodeOptions from litellm.harness.types import ( @@ -33,6 +47,7 @@ from litellm.harness.types import ( ) from litellm.llms.base_llm.harness.transformation import ( BaseCLIHarnessConfig, + HarnessSessionCopy, HarnessSessionSetup, HarnessTurnError, HarnessTurnRequest, @@ -57,6 +72,10 @@ OPENCODE_PROVIDER_NPM: Final = "@ai-sdk/openai-compatible" # A fixed title skips opencode's extra title-generation model call on the first turn. OPENCODE_SESSION_TITLE: Final = "litellm-harness" TOKEN_FILENAME: Final = "token" +RESUME_COPY_FILENAME: Final = "resume-session.json" +ID_ORDERED_PREFIX_LENGTH: Final = 16 +ID_RANDOM_LENGTH: Final = 14 +ID_RANDOM_ALPHABET: Final = string.ascii_letters + string.digits INSTRUCTIONS_FILENAME: Final = "instructions.md" XDG_DIRNAME: Final = "xdg" XDG_SUBDIRS: Final = ("config", "data", "state", "cache") @@ -162,6 +181,94 @@ class OpenCodeStreamState: step_texts: Sequence[str] = () +class ExportedPart(BaseModel): + model_config = ConfigDict(frozen=True) + + id: str + + +class ExportedMessageInfo(BaseModel): + model_config = ConfigDict(frozen=True) + + id: str + + +class ExportedMessage(BaseModel): + model_config = ConfigDict(frozen=True) + + info: ExportedMessageInfo + parts: tuple[ExportedPart, ...] + + +class ExportedSessionInfo(BaseModel): + model_config = ConfigDict(frozen=True) + + id: str + directory: str + + +class ExportedSession(BaseModel): + """The ids and folder of an `opencode export`; every other field is copied as is.""" + + model_config = ConfigDict(frozen=True) + + info: ExportedSessionInfo + messages: tuple[ExportedMessage, ...] + + def ids(self) -> Iterator[str]: + yield self.info.id + for message in self.messages: + yield message.info.id + yield from (part.id for part in message.parts) + + +EXPORT_JSON: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue) + + +def fresh_id(old: str) -> str: + """Keeps the type prefix and time-ordered part of an opencode id, so the copy sorts like the original.""" + suffix: Final = "".join(secrets.choice(ID_RANDOM_ALPHABET) for _ in range(ID_RANDOM_LENGTH)) + return f"{old[:ID_ORDERED_PREFIX_LENGTH]}{suffix}" + + +def with_id_fields_replaced(value: JsonValue, fresh: Mapping[str, str], key: str = "") -> JsonValue: + match value: + case str() if key == "id" or key.endswith("ID"): + return fresh.get(value, value) + case list(): + return [with_id_fields_replaced(item, fresh, key) for item in value] + case dict(): + return {name: with_id_fields_replaced(item, fresh, name) for name, item in value.items()} + case _: + return value + + +def filed_under(exported: JsonValue, workdir: str) -> JsonValue: + match exported: + case {"info": dict() as info}: + return {**exported, "info": {**info, "directory": workdir}} + case _: + return exported + + +def session_copy(exported: str, path: str, workdir: str) -> HarnessSessionCopy | None: + """A copy of an exported session filed under workdir, or None when the session is already filed there.""" + try: + raw: Final = EXPORT_JSON.validate_json(exported) + session: Final = ExportedSession.model_validate(raw) + except ValidationError as exc: + raise HarnessError(f"opencode export did not print a session: {exc.errors()[0]['msg']}") from exc + if session.info.directory == workdir: + return None + fresh: Final = MappingProxyType({old: fresh_id(old) for old in session.ids()}) + return HarnessSessionCopy( + session_id=fresh[session.info.id], + path=path, + data=json.dumps(filed_under(with_id_fields_replaced(raw, fresh), workdir)).encode("utf-8"), + import_command=(OPENCODE_BINARY, "--pure", "import", path), + ) + + def _as_dict(value: object) -> Mapping[str, Any]: return value if isinstance(value, dict) else MappingProxyType({}) @@ -362,6 +469,12 @@ class OpenCodeHarnessConfig(BaseCLIHarnessConfig): # The prompt goes on stdin; opencode appends non-TTY stdin to the message. return HarnessTurnRequest(argv=argv, env=setup.env, stdin=turn_prompt(ctx, prompt), cwd=ctx.sandbox.workdir) + def transform_resume_export_command(self, native_session_id: str) -> Sequence[str]: + return (OPENCODE_BINARY, "--pure", "export", native_session_id) + + def transform_resume_copy(self, exported: str, private_dir: str, workdir: str) -> HarnessSessionCopy | None: + return session_copy(exported, f"{private_dir}/{RESUME_COPY_FILENAME}", workdir) + def create_stream_state(self) -> OpenCodeStreamState: return OpenCodeStreamState() diff --git a/tests/code_coverage_tests/recursive_detector.py b/tests/code_coverage_tests/recursive_detector.py index 487c00f735a..7610d602fff 100644 --- a/tests/code_coverage_tests/recursive_detector.py +++ b/tests/code_coverage_tests/recursive_detector.py @@ -77,6 +77,7 @@ IGNORE_FUNCTIONS = [ "_restore_redacted_nested_value", # max depth set (default 10), mirrors _redact_agent_params_tree on the write side. "_unqualified", # bounded by the qualifier depth of a static TypedDict annotation (Annotated, Required/NotRequired, ReadOnly around one type, no cycles possible). "_render_json", # bounded by the nesting depth of a pydantic-validated JsonValue from the operator's config (a finite JSON tree, no cycles possible). + "with_id_fields_replaced", # harness/opencode: bounded by the nesting depth of the JsonValue that validate_json parsed from `opencode export`, which refuses nesting past ~200 levels. "completion_cost", # max depth 1: recursion only fires for mixed-tier Responses WS logging objects, and each split part carries a single service_tier so _split_responses_ws_logging_object_by_service_tier returns None. "_string_leaves", # bounded by the nesting depth of a safe_json_structure output (a finite JSON tree, no cycles possible). "_replace_string_leaves", # bounded by the nesting depth of a safe_json_structure output (a finite JSON tree, no cycles possible). diff --git a/tests/harness_e2e/test_harness_e2e.py b/tests/harness_e2e/test_harness_e2e.py index 8892085f4bc..36726c18b0f 100644 --- a/tests/harness_e2e/test_harness_e2e.py +++ b/tests/harness_e2e/test_harness_e2e.py @@ -16,7 +16,7 @@ from litellm.harness import ( ToolCall, ) -from .conftest import harness_params, model_for, requires_gateway +from .conftest import harness_available, harness_params, model_for, requires_gateway pytestmark = [requires_gateway] @@ -119,6 +119,69 @@ def test_agent_detach_and_resume(harness: Harness, workspace: Path) -> None: assert "4817" in r.text +requires_opencode = pytest.mark.skipif( + not harness_available(Harness.OPENCODE), reason="opencode runtime not installed" +) + + +@requires_opencode +def test_opencode_session_resumed_in_another_sandbox_keeps_working_in_both( + workspace: Path, tmp_path: Path +) -> None: + model = model_for(Harness.OPENCODE) + s = litellm.agent_session( + Harness.OPENCODE, + sandbox=sandbox.local(workspace), + model=model, + timeout=TURN_TIMEOUT, + ) + s.run("Remember this number: 4817. Reply with just OK.") + raw = s.detach().dumps() + moved = tmp_path / "moved" + moved.mkdir() + + with litellm.agent_resume( + State.loads(raw), sandbox=sandbox.local(workspace), model=model, timeout=TURN_TIMEOUT + ) as original: + original.run("Reply with just OK.") + with litellm.agent_resume( + State.loads(raw), sandbox=sandbox.local(moved), model=model, timeout=TURN_TIMEOUT + ) as resumed: + resumed.run( + "Create a file named hello.txt whose only content is the number I asked " + "you to remember." + ) + original.run( + "Create a file named again.txt whose only content is the number I asked " + "you to remember." + ) + + assert (moved / "hello.txt").read_text().strip() == "4817" + assert (workspace / "again.txt").read_text().strip() == "4817" + assert not (workspace / "hello.txt").exists() + + +@requires_opencode +def test_opencode_resume_of_a_missing_session_ends_with_a_runtime_error( + workspace: Path, +) -> None: + model = model_for(Harness.OPENCODE) + missing = State( + harness=Harness.OPENCODE, + native_session_id="ses_0000000000000000000missing", + workdir=str(workspace), + model=model, + ) + + with litellm.agent_resume( + missing, sandbox=sandbox.local(workspace), model=model, timeout=TURN_TIMEOUT + ) as resumed: + result = resumed.run("Reply with just OK.") + + assert result.stop_reason == "runtime_error" + assert "ses_0000000000000000000missing" in result.text + + @pytest.mark.parametrize("harness", harness_params()) def test_agent_read_only_blocks_writes(harness: Harness, workspace: Path) -> None: result = litellm.agent( diff --git a/tests/integration/sdk/test_local_sandbox_sdk.py b/tests/integration/sdk/test_local_sandbox_sdk.py new file mode 100644 index 00000000000..b30010ac3f8 --- /dev/null +++ b/tests/integration/sdk/test_local_sandbox_sdk.py @@ -0,0 +1,39 @@ +import os +import sys +from collections.abc import AsyncIterator +from pathlib import Path +from typing import Final + +import pytest +import pytest_asyncio + +from litellm import sandbox +from litellm.harness.sandbox import LocalSandbox + +WRITE_AT_PWD: Final = "import os;open(os.path.join(os.environ['PWD'],'hello.txt'),'w').write('hi')" + + +@pytest_asyncio.fixture(loop_scope="function") +async def box(tmp_path: Path) -> AsyncIterator[LocalSandbox]: + local: Final = sandbox.local(tmp_path) + yield local + await local.close() + + +@pytest.mark.asyncio(loop_scope="function") +async def test_child_trusting_pwd_writes_inside_its_cwd( + box: LocalSandbox, + monkeypatch: pytest.MonkeyPatch, + tmp_path_factory: pytest.TempPathFactory, +) -> None: + elsewhere: Final = tmp_path_factory.mktemp("elsewhere") + monkeypatch.setenv("PWD", str(elsewhere)) + os.mkdir(os.path.join(box.workdir, "sub")) + + top: Final = await box.run([sys.executable, "-I", "-c", WRITE_AT_PWD]) + nested: Final = await box.run([sys.executable, "-I", "-c", WRITE_AT_PWD], cwd="sub", env={"PWD": str(elsewhere)}) + + assert (top.exit_code, nested.exit_code) == (0, 0) + assert os.listdir(elsewhere) == [] + assert await box.read("hello.txt") == b"hi" + assert await box.read("sub/hello.txt") == b"hi" diff --git a/tests/unit/harness/sandbox/test_local.py b/tests/unit/harness/sandbox/test_local.py index 5d08dac13a1..7284a9c1157 100644 --- a/tests/unit/harness/sandbox/test_local.py +++ b/tests/unit/harness/sandbox/test_local.py @@ -1,12 +1,14 @@ +import asyncio import os import sys +from typing import Final import pytest from litellm import sandbox from litellm.harness.errors import SandboxError from litellm.harness.sandbox import LocalSandbox, Process, Sandbox -from litellm.harness.sandbox.local import filtered_environ, is_secret_env_name +from litellm.harness.sandbox.local import collect_output, filtered_environ, is_secret_env_name PY = sys.executable @@ -70,6 +72,31 @@ async def test_run_timeout_kills(sbx): await sbx.run([PY, "-c", "import time;time.sleep(30)"], timeout=0.5) +class _NeverEndingProcess: + def __init__(self) -> None: + self.stdin: asyncio.StreamWriter | None = None + self.stdout: asyncio.StreamReader = asyncio.StreamReader() + self.stderr: asyncio.StreamReader = asyncio.StreamReader() + self.killed: bool = False + + async def wait(self) -> int: + return await asyncio.Future() + + async def kill(self) -> None: + self.killed = True + + +async def test_cancelled_run_stops_its_command() -> None: + proc: Final = _NeverEndingProcess() + run: Final = asyncio.ensure_future(collect_output(proc, ["opencode", "export"], timeout=None)) + await asyncio.sleep(0) + run.cancel() + + with pytest.raises(asyncio.CancelledError): + await run + assert proc.killed + + async def test_missing_binary_raises(sbx): with pytest.raises(SandboxError): await sbx.run(["definitely-not-a-binary-xyz"]) diff --git a/tests/unit/llms/opencode/harness/test_transformation.py b/tests/unit/llms/opencode/harness/test_transformation.py index 2caf1ef80b8..c008bd28972 100644 --- a/tests/unit/llms/opencode/harness/test_transformation.py +++ b/tests/unit/llms/opencode/harness/test_transformation.py @@ -1,7 +1,9 @@ import asyncio import json +from collections.abc import Sequence from dataclasses import dataclass, field from pathlib import Path +from typing import Final import pytest from pydantic import BaseModel @@ -12,6 +14,7 @@ from litellm.harness.errors import ( HarnessError, HarnessInstallFailed, OptionsMismatch, + SandboxError, ) from litellm.harness.handlers.cli_handler import PERSIST_DIR_SCRIPT, CLIHarnessHandler from litellm.harness.options import CodexOptions, OpenCodeOptions @@ -98,6 +101,12 @@ class FakeProcess: self.killed = True +async def stall_until_timeout(cmd: Sequence[str], timeout: float | None) -> CompletedRun: + if timeout is None: + await asyncio.Event().wait() + raise SandboxError(f"command timed out after {timeout}s: {cmd[0]}") + + @dataclass class FakeSandbox: workdir: str = "/work" @@ -106,7 +115,11 @@ class FakeSandbox: outputs: list = field(default_factory=list) files: dict = field(default_factory=dict) execs: list = field(default_factory=list) - runs: list = field(default_factory=list) + stall_runs: bool = False + failing_subcommand: str | None = None + run_stdout: str = "" + runs: list[Sequence[str]] = field(default_factory=list) + run_envs: list[dict[str, str]] = field(default_factory=list) tempdirs: int = 0 async def exec(self, cmd, *, env=None, cwd=None): @@ -115,9 +128,12 @@ class FakeSandbox: async def run(self, cmd, *, env=None, cwd=None, timeout=None): self.runs.append(cmd) - if self.persist_ok: - return CompletedRun("", "", 0) - return CompletedRun("", "read-only fs", 1) + self.run_envs.append(dict(env or {})) + if self.stall_runs: + return await stall_until_timeout(cmd, timeout) + if not self.persist_ok or self.failing_subcommand in cmd: + return CompletedRun("", "read-only fs", 1) + return CompletedRun(self.run_stdout, "", 0) async def read(self, path): return self.files[path] @@ -594,13 +610,185 @@ async def test_prompt_is_sent_on_stdin_not_argv(): assert "secret prompt text" not in sandbox.execs[0]["cmd"] -async def test_resume_sets_session(): +SAVED: Final = "ses_ea6d348fdffe9pu36Y2YN6YobB" +SAVED_USER_MSG: Final = "msg_1592cb71a001307YHyO7ReafRb" +SAVED_REPLY_MSG: Final = "msg_1592cb7f4001KgUDK3Sj79XDT4" +SAVED_PARTS: Final = ("prt_1592cb71c0013NJaqSsv4dNiSR", "prt_1592cbe0f001oKdE9pRy38AduM") +COPY_PATH: Final = "/tmp/oc-1/resume-session.json" + + +def exported_session(directory: str) -> str: + user_part: Final = {"id": SAVED_PARTS[0], "sessionID": SAVED, "messageID": SAVED_USER_MSG, "type": "text"} + reply_part: Final = {"id": SAVED_PARTS[1], "sessionID": SAVED, "messageID": SAVED_REPLY_MSG, "type": "text"} + user: Final = {"id": SAVED_USER_MSG, "sessionID": SAVED, "role": "user"} + reply: Final = {"id": SAVED_REPLY_MSG, "sessionID": SAVED, "parentID": SAVED_USER_MSG, "role": "assistant"} + return json.dumps( + { + "info": {"id": SAVED, "directory": directory, "title": "litellm-harness"}, + "messages": [ + {"info": user, "parts": [{**user_part, "text": "Remember 4817"}]}, + {"info": reply, "parts": [{**reply_part, "text": "OK"}]}, + ], + } + ) + + +async def resumed(saved_in: str) -> tuple[CLIHarnessHandler, SessionContext, FakeSandbox]: handler, ctx, sandbox = await started() - await handler.resume(ctx, "ses_prev") + sandbox.run_stdout = exported_session(saved_in) + await handler.resume(ctx, SAVED) sandbox.outputs.append(fixture_proc("turn2_session_skill.jsonl")) await collect(handler, ctx, "hi") - cmd = sandbox.execs[0]["cmd"] - assert cmd[cmd.index("--session") + 1] == "ses_prev" + return handler, ctx, sandbox + + +def resumed_session_arg(sandbox: FakeSandbox) -> str: + cmd: Final = sandbox.execs[0]["cmd"] + return cmd[cmd.index("--session") + 1] + + +async def test_resume_elsewhere_continues_a_copy_filed_under_the_sandbox() -> None: + _, _, sandbox = await resumed("/other") + + copy: Final = json.loads(sandbox.files[COPY_PATH]) + copy_id: Final = resumed_session_arg(sandbox) + assert copy["info"]["id"] == copy_id != SAVED + assert copy["info"]["directory"] == sandbox.workdir + assert tuple(sandbox.runs[-1]) == ("opencode", "--pure", "import", COPY_PATH) + assert sandbox.run_envs[-1]["XDG_DATA_HOME"] == "/tmp/oc-1/xdg/data" + assert sandbox.run_envs[-2]["XDG_DATA_HOME"] == "/tmp/oc-1/xdg/data" + + +async def test_resume_copy_gets_fresh_ids_that_keep_their_links_and_order() -> None: + _, _, sandbox = await resumed("/other") + + raw_copy: Final = sandbox.files[COPY_PATH].decode() + user, reply = json.loads(raw_copy)["messages"] + copy_id: Final = resumed_session_arg(sandbox) + assert all(old not in raw_copy for old in (SAVED, SAVED_USER_MSG, SAVED_REPLY_MSG, *SAVED_PARTS)) + assert {user["info"]["sessionID"], reply["info"]["sessionID"]} == {copy_id} + assert reply["info"]["parentID"] == user["info"]["id"] + assert user["parts"][0]["messageID"] == user["info"]["id"] + assert reply["parts"][0]["messageID"] == reply["info"]["id"] + assert {user["parts"][0]["sessionID"], reply["parts"][0]["sessionID"]} == {copy_id} + assert (user["parts"][0]["text"], reply["parts"][0]["text"]) == ("Remember 4817", "OK") + + +def test_resume_copy_keeps_the_message_order() -> None: + message_ids: Final = tuple(f"msg_1592cb71a{n:03x}307YHyO7ReafRb" for n in range(10)) + messages: Final = [{"info": {"id": message_id}, "parts": []} for message_id in message_ids] + exported: Final = json.dumps({"info": {"id": SAVED, "directory": "/other"}, "messages": messages}) + + copy: Final = CONFIG.transform_resume_copy(exported, PRIVATE, "/work") + + assert copy is not None + copied_ids: Final = [message["info"]["id"] for message in json.loads(copy.data)["messages"]] + assert sorted(copied_ids) == copied_ids + + +def test_resume_copy_keeps_saved_text_that_reads_like_an_id() -> None: + tool_part: Final = { + "id": SAVED_PARTS[1], + "sessionID": SAVED, + "messageID": SAVED_USER_MSG, + "type": "tool", + "state": {"input": {"command": SAVED_REPLY_MSG}, "output": SAVED_USER_MSG}, + } + text_part: Final = {"id": SAVED_PARTS[0], "sessionID": SAVED, "messageID": SAVED_USER_MSG, "type": "text", "text": SAVED} + message: Final = {"info": {"id": SAVED_USER_MSG, "sessionID": SAVED}, "parts": [text_part, tool_part]} + exported: Final = json.dumps({"info": {"id": SAVED, "directory": "/other"}, "messages": [message]}) + + copy: Final = CONFIG.transform_resume_copy(exported, PRIVATE, "/work") + + assert copy is not None + text, tool = json.loads(copy.data)["messages"][0]["parts"] + assert text["text"] == SAVED != text["sessionID"] == copy.session_id + assert (tool["state"]["input"]["command"], tool["state"]["output"]) == (SAVED_REPLY_MSG, SAVED_USER_MSG) + assert tool["messageID"] != SAVED_USER_MSG + + +async def test_later_turns_after_a_resume_do_not_export_again() -> None: + handler, ctx, sandbox = await resumed("/other") + + sandbox.outputs.append(fixture_proc("turn2_session_skill.jsonl")) + await collect(handler, ctx, "again") + + assert [run[2] for run in sandbox.runs if run[0] == "opencode"] == ["export", "import"] + + +async def test_each_resume_elsewhere_makes_its_own_copy() -> None: + _, _, first = await resumed("/other") + _, _, second = await resumed("/other") + + assert resumed_session_arg(first) != resumed_session_arg(second) + + +async def test_resume_in_the_folder_it_was_saved_in_continues_the_original() -> None: + _, _, sandbox = await resumed("/work") + + assert resumed_session_arg(sandbox) == SAVED + assert COPY_PATH not in sandbox.files + assert all("import" not in run for run in sandbox.runs) + + +@pytest.mark.parametrize("failure", ["export_exits_1", "import_exits_1", "export_stalls", "export_unreadable"]) +async def test_resume_that_cannot_be_prepared_fails_the_turn_without_running_opencode( + failure: str, +) -> None: + handler, ctx, sandbox = await started() + sandbox.run_stdout = "Error: not a session" if failure == "export_unreadable" else exported_session("/other") + sandbox.failing_subcommand = failure.split("_")[0] if failure.endswith("exits_1") else None + sandbox.stall_runs = failure == "export_stalls" + + await asyncio.wait_for(handler.resume(ctx, SAVED), 5) + with pytest.raises(HarnessTurnError, match=SAVED): + await collect(handler, ctx, "hi") + assert sandbox.execs == [] + + +async def test_resume_runs_nothing_until_the_first_turn() -> None: + handler, ctx, sandbox = await started() + runs_before_resume: Final = len(sandbox.runs) + + await handler.resume(ctx, SAVED) + + assert len(sandbox.runs) == runs_before_resume + assert COPY_PATH not in sandbox.files + + +async def test_restart_after_a_failed_resume_tries_the_resume_again() -> None: + handler, ctx, sandbox = await started() + sandbox.run_stdout = exported_session("/other") + sandbox.failing_subcommand = "export" + await handler.resume(ctx, SAVED) + with pytest.raises(HarnessTurnError, match=SAVED): + await collect(handler, ctx, "hi") + + sandbox.failing_subcommand = None + await handler.start(ctx) + await handler.resume(ctx, SAVED) + sandbox.outputs.append(fixture_proc("turn2_session_skill.jsonl")) + await collect(handler, ctx, "again") + + assert resumed_session_arg(sandbox) == json.loads(sandbox.files["/tmp/oc-2/resume-session.json"])["info"]["id"] + + +async def test_restart_resumes_its_own_session_without_exporting_it() -> None: + handler, ctx, sandbox = await started() + sandbox.outputs.append(fixture_proc("turn1_write_read.jsonl")) + await collect(handler, ctx, "first") + native_id: Final = handler.native_session_id() + assert native_id is not None + runs_before_restart: Final = len(sandbox.runs) + + await handler.start(ctx) + await handler.resume(ctx, native_id) + sandbox.outputs.append(fixture_proc("turn2_session_skill.jsonl")) + await collect(handler, ctx, "again") + + assert all("export" not in run for run in sandbox.runs[runs_before_restart:]) + cmd: Final = sandbox.execs[1]["cmd"] + assert cmd[cmd.index("--session") + 1] == native_id async def test_read_only_and_disable_tools_config():