a2a/tests/test_workspace.py

from __future__ import annotations

import pytest
from pydantic import BaseModel

from a2a_pack import (
    A2AAgent,
    FileType,
    LocalWorkspaceClient,
    NoAuth,
    RunContext,
    WorkspaceAccess,
    WorkspaceDenied,
    WorkspaceMode,
    skill,
)


_FILES: dict[str, bytes] = {
    "src/auth/login.py": b"def login(jwt): ...  # JWT auth middleware",
    "src/auth/middleware.py": b"# auth middleware for JWT validation",
    "src/payments/checkout.py": b"def checkout(): ...  # payment flow",
    "tests/test_auth.py": b"def test_login_jwt(): ...",
    "configs/app.toml": b"[auth]\njwt = true",
    "secrets/.env": b"DB_PASSWORD=oops",
    "README.md": b"# project",
}


class _Cfg(BaseModel):
    pass


class _CoderAgent(A2AAgent[_Cfg, NoAuth]):
    name = "coder"
    description = "Edits code by negotiated views"
    workspace_access = WorkspaceAccess.dynamic(
        max_files=5,
        allowed_modes=(
            WorkspaceMode.READ_ONLY,
            WorkspaceMode.READ_WRITE_OVERLAY,
        ),
        require_reason=True,
        deny_patterns=("secrets/**", ".env", "**/.env"),
    )

    @skill()
    async def find_and_patch_auth(self, ctx: RunContext[NoAuth]) -> int:
        view = await ctx.workspace.open_view(
            purpose="Fix JWT login bug",
            hints=["auth", "jwt", "login"],
            file_types=[FileType.PYTHON],
            max_files=3,
            mode=WorkspaceMode.READ_WRITE_OVERLAY,
        )
        for fm in view.files:
            content = await view.read(fm.path)
            await view.write(fm.path, content + b"\n# patched\n")
        return len(view.files)


def _client() -> LocalWorkspaceClient:
    return LocalWorkspaceClient(_FILES, access=_CoderAgent.workspace_access)


async def test_open_view_grants_relevant_files_only():
    agent = _CoderAgent()
    n = await agent.local_invoke("find_and_patch_auth", workspace=_client())
    assert n >= 1
    assert n <= 3


async def test_search_excludes_denied_patterns():
    ws = _client()
    matches = await ws.search(query="DB_PASSWORD env secret", limit=20)
    assert all("secrets/" not in m.path for m in matches)
    assert all(not m.path.endswith(".env") for m in matches)


async def test_request_access_rejects_denied_path():
    ws = _client()
    with pytest.raises(WorkspaceDenied, match="denied by policy"):
        await ws.request_access(
            files=["secrets/.env"],
            mode=WorkspaceMode.READ_ONLY,
            reason="trying to read secrets",
        )


async def test_request_access_rejects_disallowed_mode():
    ws = _client()
    with pytest.raises(WorkspaceDenied, match="not in allowed_modes"):
        await ws.request_access(
            files=["src/auth/login.py"],
            mode=WorkspaceMode.READ_WRITE_DIRECT,
            reason="needs direct write",
        )


async def test_request_access_requires_reason():
    ws = _client()
    with pytest.raises(WorkspaceDenied, match="reason required"):
        await ws.request_access(
            files=["src/auth/login.py"],
            mode=WorkspaceMode.READ_ONLY,
            reason="",
        )


async def test_request_access_enforces_max_files():
    ws = _client()
    paths = [p for p in _FILES if not p.startswith("secrets") and not p.endswith(".env")]
    assert len(paths) > 5  # confirm fixture
    with pytest.raises(WorkspaceDenied, match="max_files"):
        await ws.request_access(
            files=paths[:6],
            mode=WorkspaceMode.READ_ONLY,
            reason="too many",
        )


async def test_writes_are_staged_as_patches_not_applied():
    ws = _client()
    grant = await ws.request_access(
        files=["src/auth/login.py"],
        mode=WorkspaceMode.READ_WRITE_OVERLAY,
        reason="patch",
    )
    from a2a_pack.workspace import LocalWorkspaceView

    view = LocalWorkspaceView(grant, ws)
    patch = await view.write("src/auth/login.py", b"new content")
    assert patch.operation == "update"
    assert patch.content == b"new content"
    # original file untouched in the in-memory store
    assert ws._files["src/auth/login.py"].startswith(b"def login")
    patches = await view.patches()
    assert len(patches) == 1


async def test_view_rejects_writes_in_read_only_mode():
    ws = _client()
    grant = await ws.request_access(
        files=["src/auth/login.py"],
        mode=WorkspaceMode.READ_ONLY,
        reason="reading",
    )
    from a2a_pack.workspace import LocalWorkspaceView

    view = LocalWorkspaceView(grant, ws)
    with pytest.raises(WorkspaceDenied, match="read-only"):
        await view.write("src/auth/login.py", b"x")


async def test_view_rejects_path_outside_grant():
    ws = _client()
    grant = await ws.request_access(
        files=["src/auth/login.py"],
        mode=WorkspaceMode.READ_WRITE_OVERLAY,
        reason="patching",
    )
    from a2a_pack.workspace import LocalWorkspaceView

    view = LocalWorkspaceView(grant, ws)
    with pytest.raises(WorkspaceDenied, match="not in grant"):
        await view.read("src/payments/checkout.py")


async def test_workspace_disabled_by_default():
    from a2a_pack import SkillInvocationError

    class _Plain(A2AAgent):
        name = "plain"
        description = ""

        @skill()
        async def touch(self, ctx: RunContext[NoAuth]) -> str:
            await ctx.workspace.search(query="x")
            return "ok"

    agent = _Plain()
    with pytest.raises(SkillInvocationError) as ei:
        await agent.local_invoke("touch")
    assert isinstance(ei.value.__cause__, PermissionError)


def test_workspace_access_propagates_to_card():
    card = _CoderAgent().card()
    wa = card.workspace_access
    assert wa.enabled is True
    assert wa.max_files == 5
    assert WorkspaceMode.READ_WRITE_OVERLAY in wa.allowed_modes
    assert "secrets/**" in wa.deny_patterns


def test_default_sandbox_is_microsandbox():
    from a2a_pack import Sandbox

    assert _CoderAgent.runtime().sandbox is Sandbox.MICROSANDBOX