dify/dify-agent/tests/local/dify_agent/adapters/shell/test_shellctl.py

"""Local tests for the shellctl shell adapter and env-driven provider factory."""

from __future__ import annotations

import asyncio
import base64
from collections.abc import Callable
from dataclasses import dataclass, field
from typing import cast

import httpx2 as httpx
import pytest
from pydantic import ValidationError

from dify_agent.adapters.shell import shellctl
from dify_agent.adapters.shell.config import ShellAdapterSettings
from dify_agent.adapters.shell.factory import create_shell_provider
from dify_agent.adapters.shell.protocols import ShellCommandResult, ShellProviderError
from dify_agent.adapters.shell.shellctl import (
    ShellctlClientProtocol,
    ShellFileTransferError,
    ShellctlProvider,
)


@dataclass(slots=True)
class _Job:
    job_id: str
    status: str = "running"
    done: bool = True
    output: str = ""
    offset: int = 0
    truncated: bool = False
    exit_code: int | None = 0
    output_path: str | None = "/tmp/output.log"


@dataclass(slots=True)
class _Status:
    job_id: str
    status: str = "terminated"
    done: bool = True
    offset: int = 0
    exit_code: int | None = 130


@dataclass(slots=True)
class _RunCall:
    script: str
    cwd: str | None
    env: dict[str, str] | None
    timeout: float


type _RunHandler = Callable[[str, str | None, dict[str, str] | None, float], _Job]
type _WaitHandler = Callable[[str, int, float], _Job]
type _InputHandler = Callable[[str, str, int, float], _Job]
type _TerminateHandler = Callable[[str, float], _Status]


@dataclass(slots=True)
class FakeShellctlClient:
    run_handler: _RunHandler | None = None
    wait_handler: _WaitHandler | None = None
    input_handler: _InputHandler | None = None
    tail_handler: Callable[[str], _Job] | None = None
    terminate_handler: _TerminateHandler | None = None
    run_calls: list[_RunCall] = field(default_factory=list)
    wait_calls: list[tuple[str, int, float]] = field(default_factory=list)
    input_calls: list[tuple[str, str, int, float]] = field(default_factory=list)
    terminate_calls: list[tuple[str, float]] = field(default_factory=list)
    delete_calls: list[tuple[str, bool, float | None]] = field(default_factory=list)
    closed: bool = False

    async def run(
        self,
        script: str,
        *,
        cwd: str | None = None,
        env: dict[str, str] | None = None,
        timeout: float = 30.0,
    ) -> _Job:
        self.run_calls.append(_RunCall(script=script, cwd=cwd, env=env, timeout=timeout))
        if self.run_handler is not None:
            return self.run_handler(script, cwd, env, timeout)
        return _Job(job_id="job", status="exited", done=True, exit_code=0)

    async def wait(self, job_id: str, *, offset: int, timeout: float = 30.0) -> _Job:
        self.wait_calls.append((job_id, offset, timeout))
        if self.wait_handler is not None:
            return self.wait_handler(job_id, offset, timeout)
        return _Job(job_id=job_id, status="exited", done=True, offset=offset, exit_code=0)

    async def input(
        self,
        job_id: str,
        text: str,
        *,
        offset: int,
        timeout: float = 30.0,
    ) -> _Job:
        self.input_calls.append((job_id, text, offset, timeout))
        if self.input_handler is not None:
            return self.input_handler(job_id, text, offset, timeout)
        return _Job(job_id=job_id, status="exited", done=True, offset=offset, exit_code=0)

    async def tail(self, job_id: str) -> _Job:
        if self.tail_handler is not None:
            return self.tail_handler(job_id)
        return _Job(job_id=job_id, status="exited", done=True, output="", exit_code=0)

    async def terminate(self, job_id: str, grace_seconds: float = 10.0) -> _Status:
        self.terminate_calls.append((job_id, grace_seconds))
        if self.terminate_handler is not None:
            return self.terminate_handler(job_id, grace_seconds)
        return _Status(job_id=job_id)

    async def delete(
        self,
        job_id: str,
        *,
        force: bool = False,
        grace_seconds: float | None = None,
    ) -> None:
        self.delete_calls.append((job_id, force, grace_seconds))
        return None

    async def close(self) -> None:
        self.closed = True


def _provider(client: FakeShellctlClient) -> ShellctlProvider:
    return ShellctlProvider(
        entrypoint="http://shellctl",
        token="",
        client_factory=lambda: _client_protocol(client),
    )


def _client_protocol(client: FakeShellctlClient) -> ShellctlClientProtocol:
    return cast(ShellctlClientProtocol, cast(object, client))


def test_factory_unknown_provider_raises() -> None:
    with pytest.raises(ValidationError):
        ShellAdapterSettings(shell_provider="nope")  # type: ignore[arg-type]


def test_factory_shellctl_requires_entrypoint() -> None:
    with pytest.raises(ValidationError, match="shellctl_entrypoint is required"):
        ShellAdapterSettings(shell_provider="shellctl", shellctl_entrypoint=None)


def test_factory_builds_shellctl_provider_from_settings() -> None:
    settings = ShellAdapterSettings(shell_provider="shellctl", shellctl_entrypoint="http://shellctl.example")
    provider = create_shell_provider(settings)
    assert isinstance(provider, ShellctlProvider)
    assert provider.entrypoint == "http://shellctl.example"
    assert provider.token == ""


def test_provider_create_opens_only_live_resource_and_suspend_closes_client() -> None:
    client = FakeShellctlClient()

    async def scenario() -> None:
        resource = await _provider(client).create()
        assert client.run_calls == []
        await resource.suspend()

    asyncio.run(scenario())
    assert client.closed is True


def test_commands_forward_parameters_and_map_metadata() -> None:
    client = FakeShellctlClient(
        run_handler=lambda script, cwd, env, timeout: _Job(
            job_id="run-job",
            status="running",
            done=False,
            output="abc",
            offset=3,
            truncated=True,
            exit_code=None,
            output_path="/tmp/run.log",
        ),
        wait_handler=lambda job_id, offset, timeout: _Job(
            job_id=job_id,
            status="running",
            done=False,
            output="def",
            offset=6,
            truncated=False,
            exit_code=None,
            output_path="/tmp/run.log",
        ),
        input_handler=lambda job_id, text, offset, timeout: _Job(
            job_id=job_id,
            status="exited",
            done=True,
            output="ghi",
            offset=9,
            truncated=False,
            exit_code=0,
            output_path="/tmp/run.log",
        ),
        tail_handler=lambda job_id: _Job(
            job_id=job_id,
            status="exited",
            done=True,
            output="tail",
            offset=11,
            truncated=False,
            exit_code=0,
            output_path="/tmp/tail.log",
        ),
        terminate_handler=lambda job_id, grace_seconds: _Status(
            job_id=job_id,
            status="terminated",
            done=True,
            offset=12,
            exit_code=130,
        ),
    )

    async def scenario() -> None:
        resource = await _provider(client).create()
        run_result = await resource.commands.run("pwd", cwd="~/workspace/abc12ff", env={"FOO": "bar"}, timeout=2.5)
        wait_result = await resource.commands.wait("run-job", offset=3, timeout=4.0)
        read_result = await resource.commands.read_output("run-job", offset=6)
        input_result = await resource.commands.input("run-job", "ls\n", offset=6, timeout=5.0)
        interrupt_result = await resource.commands.interrupt("run-job", grace_seconds=1.5)
        tail_result = await resource.commands.tail("run-job")
        await resource.commands.delete("run-job", force=True, grace_seconds=2.0)
        await resource.suspend()

        assert run_result == ShellCommandResult(
            job_id="run-job",
            status="running",
            done=False,
            exit_code=None,
            output="abc",
            offset=3,
            truncated=True,
            output_path="/tmp/run.log",
        )
        assert wait_result.offset == 6
        assert read_result.offset == 6
        assert input_result.exit_code == 0
        assert interrupt_result.status == "terminated"
        assert tail_result.output_path == "/tmp/tail.log"

    asyncio.run(scenario())

    assert client.run_calls == [_RunCall(script="pwd", cwd="~/workspace/abc12ff", env={"FOO": "bar"}, timeout=2.5)]
    assert client.wait_calls == [
        ("run-job", 3, 4.0),
        ("run-job", 6, 0.0),
    ]
    assert client.input_calls == [("run-job", "ls\n", 6, 5.0)]
    assert client.terminate_calls == [("run-job", 1.5)]
    assert client.delete_calls == [("run-job", True, 2.0)]


def test_commands_map_http_timeout_to_shell_provider_error() -> None:
    request = httpx.Request("POST", "http://shellctl.example/v1/jobs")
    client = FakeShellctlClient(
        run_handler=lambda script, cwd, env, timeout: (_ for _ in ()).throw(
            httpx.ReadTimeout("timed out", request=request)
        )
    )

    async def scenario() -> None:
        resource = await _provider(client).create()
        with pytest.raises(ShellProviderError, match="timed out") as exc_info:
            await resource.commands.run("pwd", timeout=2.5)
        assert exc_info.value.code == "timeout"

    asyncio.run(scenario())


def test_commands_map_http_request_error_to_shell_provider_error() -> None:
    request = httpx.Request("POST", "http://shellctl.example/v1/jobs/run")
    client = FakeShellctlClient(
        wait_handler=lambda job_id, offset, timeout: (_ for _ in ()).throw(
            httpx.ConnectError("connection failed", request=request)
        )
    )

    async def scenario() -> None:
        resource = await _provider(client).create()
        with pytest.raises(ShellProviderError, match="connection failed") as exc_info:
            await resource.commands.wait("run-job", offset=3, timeout=4.0)
        assert exc_info.value.code == "request_error"

    asyncio.run(scenario())


def test_delete_maps_http_timeout_to_shell_provider_error() -> None:
    request = httpx.Request("DELETE", "http://shellctl.example/v1/jobs/run-job")

    @dataclass(slots=True)
    class DeleteTimeoutClient(FakeShellctlClient):
        async def delete(self, job_id, *, force=False, grace_seconds=None):
            self.delete_calls.append((job_id, force, grace_seconds))
            raise httpx.ReadTimeout("delete timed out", request=request)

    client = DeleteTimeoutClient()

    async def scenario() -> None:
        resource = await _provider(client).create()
        with pytest.raises(ShellProviderError, match="delete timed out") as exc_info:
            await resource.commands.delete("run-job", force=True, grace_seconds=2.0)
        assert exc_info.value.code == "timeout"

    asyncio.run(scenario())
    assert client.delete_calls == [("run-job", True, 2.0)]


def test_delete_maps_http_request_error_to_shell_provider_error() -> None:
    request = httpx.Request("DELETE", "http://shellctl.example/v1/jobs/run-job")

    @dataclass(slots=True)
    class DeleteRequestErrorClient(FakeShellctlClient):
        async def delete(self, job_id, *, force=False, grace_seconds=None):
            self.delete_calls.append((job_id, force, grace_seconds))
            raise httpx.ConnectError("delete connection failed", request=request)

    client = DeleteRequestErrorClient()

    async def scenario() -> None:
        resource = await _provider(client).create()
        with pytest.raises(ShellProviderError, match="delete connection failed") as exc_info:
            await resource.commands.delete("run-job", force=True, grace_seconds=2.0)
        assert exc_info.value.code == "request_error"

    asyncio.run(scenario())
    assert client.delete_calls == [("run-job", True, 2.0)]


def test_files_upload_and_download_still_work() -> None:
    content = b"hello \x00 world"
    encoded = base64.b64encode(content).decode("ascii")
    client = FakeShellctlClient(
        run_handler=lambda script, cwd, env, timeout: (
            _Job(job_id="ul-job", status="exited", done=True, exit_code=0)
            if "base64 -d" in script
            else _Job(
                job_id="dl-job",
                status="exited",
                done=True,
                exit_code=0,
                output=f"noise{shellctl._TRANSFER_BEGIN}{encoded}{shellctl._TRANSFER_END}tail",
            )
        )
    )

    async def scenario() -> None:
        resource = await _provider(client).create()
        await resource.files.upload(content=content, remote_path="out.bin", cwd="~/workspace/abc12ff")
        downloaded = await resource.files.download(remote_path="report.txt", cwd="~/workspace/abc12ff")
        assert downloaded == content

    asyncio.run(scenario())


def test_file_transfer_timeout_is_an_end_to_end_budget(monkeypatch: pytest.MonkeyPatch) -> None:
    clock = {"value": 100.0}

    def fake_monotonic() -> float:
        return clock["value"]

    monkeypatch.setattr(shellctl.time, "monotonic", fake_monotonic)

    def run_handler(script: str, cwd: str | None, env: dict[str, str] | None, timeout: float) -> _Job:
        del script, cwd, env
        assert timeout == pytest.approx(5.0, rel=0, abs=0.01)
        clock["value"] = 103.5
        return _Job(job_id="upload-job", status="running", done=False, output="part-1", offset=6, exit_code=None)

    def wait_handler(job_id: str, offset: int, timeout: float) -> _Job:
        assert job_id == "upload-job"
        assert offset == 6
        assert timeout == pytest.approx(1.5, rel=0, abs=0.01)
        return _Job(job_id=job_id, status="exited", done=True, output="part-2", offset=12, exit_code=0)

    client = FakeShellctlClient(run_handler=run_handler, wait_handler=wait_handler)

    async def scenario() -> None:
        transfer = shellctl.ShellctlFileTransfer(
            client=_client_protocol(client),
            timeout=5.0,
        )
        await transfer.upload(content=b"payload", remote_path="out.bin")

    asyncio.run(scenario())
    assert client.delete_calls == [("upload-job", True, None)]


def test_file_transfer_timeout_exhaustion_raises_timeout_and_still_deletes_job(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    clock = {"value": 100.0}

    def fake_monotonic() -> float:
        return clock["value"]

    monkeypatch.setattr(shellctl.time, "monotonic", fake_monotonic)

    def run_handler(script: str, cwd: str | None, env: dict[str, str] | None, timeout: float) -> _Job:
        del script, cwd, env
        assert timeout == pytest.approx(5.0, rel=0, abs=0.01)
        clock["value"] = 106.0
        return _Job(job_id="upload-job", status="running", done=False, output="part-1", offset=6, exit_code=None)

    client = FakeShellctlClient(run_handler=run_handler)

    async def scenario() -> None:
        transfer = shellctl.ShellctlFileTransfer(
            client=_client_protocol(client),
            timeout=5.0,
        )
        with pytest.raises(ShellProviderError, match="timed out") as exc_info:
            await transfer.upload(content=b"payload", remote_path="out.bin")
        assert exc_info.value.code == "timeout"

    asyncio.run(scenario())
    assert client.delete_calls == [("upload-job", True, None)]


def test_download_missing_file_raises() -> None:
    client = FakeShellctlClient(
        run_handler=lambda script, cwd, env, timeout: _Job(
            job_id="dl-job",
            status="exited",
            done=True,
            output="",
            exit_code=shellctl._DOWNLOAD_MISSING_EXIT_CODE,
        )
    )

    async def scenario() -> None:
        resource = await _provider(client).create()
        with pytest.raises(ShellFileTransferError, match="not found"):
            await resource.files.download(remote_path="missing.txt")

    asyncio.run(scenario())