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 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:
    settings = ShellAdapterSettings(shell_provider="nope")
    with pytest.raises(ValueError, match="Unknown shell provider"):
        create_shell_provider(settings)


def test_factory_shellctl_requires_entrypoint() -> None:
    settings = ShellAdapterSettings(shell_provider="shellctl", shellctl_entrypoint=None)
    with pytest.raises(ValueError, match="DIFY_AGENT_SHELLCTL_ENTRYPOINT"):
        create_shell_provider(settings)


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_close_closes_client() -> None:
    client = FakeShellctlClient()

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

    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.close()

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