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():