strukto-ai--mirage/python/mirage/agents/pydantic_ai/backend.py

# ========= Copyright 2026 @ Strukto.AI All Rights Reserved. =========
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ========= Copyright 2026 @ Strukto.AI All Rights Reserved. =========

import shlex

from pydantic_ai import BinaryContent, ToolReturn
from pydantic_ai_backends.protocol import SandboxProtocol
from pydantic_ai_backends.types import (EditResult, ExecuteResponse, FileInfo,
                                        GrepMatch, WriteResult)

from mirage.agents.pydantic_ai._convert import (io_to_execute_response,
                                                io_to_file_infos,
                                                io_to_grep_matches)
from mirage.bridge.sync import run_async_from_sync
from mirage.core.filetype.pdf import pages_to_images
from mirage.io.types import IOResult
from mirage.types import DEFAULT_SESSION_ID
from mirage.workspace.workspace import Workspace


class PydanticAIWorkspace(SandboxProtocol):
    """Pydantic AI backend backed by a Mirage Workspace.

    File operations (read, write, edit, ls) go through the Ops layer directly.
    Shell operations (execute, grep, glob) go through Workspace.execute()
    for pipe and flag support.
    """

    def __init__(
        self,
        workspace: Workspace,
        sandbox_id: str = "mirage",
        session_id: str = DEFAULT_SESSION_ID,
    ) -> None:
        self._ws = workspace
        self._id = sandbox_id
        self._session_id = session_id

    def _run(self, coro):
        return run_async_from_sync(coro)

    @property
    def id(self) -> str:
        return self._id

    async def _exec(self, command: str) -> IOResult:
        return await self._ws.execute(command, session_id=self._session_id)

    def _read_bytes(self, path: str) -> bytes:
        return self._run(self._aread_bytes(path))

    async def _aread_bytes(self, path: str) -> bytes:
        ops = self._ws.ops
        return await ops.read(path)

    # -- execute -------------------------------------------------------

    def execute(self,
                command: str,
                timeout: int | None = None) -> ExecuteResponse:
        return self._run(self.aexecute(command, timeout=timeout))

    async def aexecute(self,
                       command: str,
                       timeout: int | None = None) -> ExecuteResponse:
        io = await self._exec(command)
        return io_to_execute_response(io)

    # -- ls_info -------------------------------------------------------

    def ls_info(self, path: str) -> list[FileInfo]:
        return self._run(self.als_info(path))

    async def als_info(self, path: str) -> list[FileInfo]:
        io = await self._exec(f"ls {shlex.quote(path)}")
        stdout = (io.stdout or b"").decode("utf-8", errors="replace").strip()
        if not stdout:
            return []
        base = path.rstrip("/")
        result: list[FileInfo] = []
        for name in stdout.split("\n"):
            name = name.strip()
            if not name:
                continue
            is_dir = name.endswith("/")
            clean = name.rstrip("/")
            result.append(
                FileInfo(name=clean,
                         path=f"{base}/{clean}",
                         is_dir=is_dir,
                         size=None))
        return result

    # -- read ----------------------------------------------------------

    def read(self,
             path: str,
             offset: int = 0,
             limit: int = 2000) -> str | ToolReturn:
        return self._run(self.aread(path, offset, limit))

    async def aread(self,
                    path: str,
                    offset: int = 0,
                    limit: int = 2000) -> str | ToolReturn:
        ops = self._ws.ops
        if path.lower().endswith(".pdf"):
            return await self._aread_pdf(path)
        try:
            data = await ops.read(path)
        except (FileNotFoundError, ValueError) as exc:
            return f"Error: {exc}"
        text = data.decode("utf-8", errors="replace")
        lines = text.splitlines(keepends=True)
        sliced = lines[offset:offset + limit]
        numbered = []
        for i, line in enumerate(sliced, start=offset + 1):
            numbered.append(f"{i:>6}\t{line}")
        return "".join(numbered)

    async def _aread_pdf(self, path: str) -> str | ToolReturn:
        ops = self._ws.ops
        try:
            raw = await ops.read(path)
        except (FileNotFoundError, ValueError) as exc:
            return f"Error: {exc}"
        images = pages_to_images(raw)
        filename = path.rsplit("/", 1)[-1]
        content: list[str | BinaryContent] = []
        for page_num, png_bytes in images:
            content.append(f"Page {page_num}:")
            content.append(
                BinaryContent(data=png_bytes, media_type="image/png"))
        return ToolReturn(
            return_value=f"PDF: {filename} ({len(images)} pages)",
            content=content,
        )

    # -- write ---------------------------------------------------------

    def write(self, path: str, content: str | bytes) -> WriteResult:
        return self._run(self.awrite(path, content))

    async def awrite(self, path: str, content: str | bytes) -> WriteResult:
        ops = self._ws.ops
        try:
            await ops.stat(path)
            return WriteResult(error=f"Error: file '{path}' already exists")
        except (FileNotFoundError, ValueError):
            pass
        parent = "/".join(path.rstrip("/").split("/")[:-1]) or "/"
        try:
            await ops.mkdir(parent)
        except (FileExistsError, ValueError):
            pass
        data = content.encode("utf-8") if isinstance(content, str) else content
        await ops.write(path, data)
        return WriteResult(path=path)

    # -- edit ----------------------------------------------------------

    def edit(
        self,
        path: str,
        old_string: str,
        new_string: str,
        replace_all: bool = False,
    ) -> EditResult:
        return self._run(self.aedit(path, old_string, new_string, replace_all))

    async def aedit(
        self,
        path: str,
        old_string: str,
        new_string: str,
        replace_all: bool = False,
    ) -> EditResult:
        ops = self._ws.ops
        try:
            data = await ops.read(path)
        except (FileNotFoundError, ValueError):
            return EditResult(error=f"Error: file '{path}' not found")
        content = data.decode("utf-8", errors="replace")
        count = content.count(old_string)
        if count == 0:
            return EditResult(
                error=f"Error: string not found in file: '{old_string}'")
        if count > 1 and not replace_all:
            return EditResult(
                error=f"Error: string '{old_string}' appears {count} times. "
                f"Use replace_all=True")
        if replace_all:
            new_content = content.replace(old_string, new_string)
        else:
            new_content = content.replace(old_string, new_string, 1)
        await ops.write(path, new_content.encode("utf-8"))
        return EditResult(path=path, occurrences=count if replace_all else 1)

    # -- grep_raw ------------------------------------------------------

    def grep_raw(
        self,
        pattern: str,
        path: str | None = None,
        glob: str | None = None,
        ignore_hidden: bool = True,
    ) -> list[GrepMatch] | str:
        return self._run(self.agrep_raw(pattern, path, glob, ignore_hidden))

    async def agrep_raw(
        self,
        pattern: str,
        path: str | None = None,
        glob: str | None = None,
        ignore_hidden: bool = True,
    ) -> list[GrepMatch] | str:
        parts = ["grep", "-rn"]
        if glob:
            parts.extend(["--include", shlex.quote(glob)])
        parts.append(shlex.quote(pattern))
        parts.append(shlex.quote(path or "/"))
        io = await self._exec(" ".join(parts))
        return io_to_grep_matches(io)

    # -- glob_info -----------------------------------------------------

    def glob_info(self, pattern: str, path: str = "/") -> list[FileInfo]:
        return self._run(self.aglob_info(pattern, path))

    async def aglob_info(self,
                         pattern: str,
                         path: str = "/") -> list[FileInfo]:
        name = pattern.split("/")[-1] if "/" in pattern else pattern
        io = await self._exec(
            f"find {shlex.quote(path)} -name {shlex.quote(name)}")
        return io_to_file_infos(io)