diff --git a/sdk/python/README.md b/sdk/python/README.md index a3f854e..cdc4d8c 100644 --- a/sdk/python/README.md +++ b/sdk/python/README.md @@ -10,6 +10,12 @@ isolated AFS-backed workspace. pip install redis-afs ``` +For the async client, install the optional `async` extra (installs `httpx`): + +```bash +pip install 'redis-afs[async]' +``` + ## Quick Start ```python @@ -39,6 +45,51 @@ with afs.fs.mount(workspaces=[{"name": "foobar"}], mode="rw") as fs: fs.write_file("/README.md", "hello") ``` +## Async usage + +The async client mirrors the sync surface, but every network call is a +coroutine. Install the `async` extra (`pip install 'redis-afs[async]'`) and +import `AsyncAFS`. + +```python +import os +from redis_afs import AsyncAFS + +async def main(): + async with AsyncAFS(api_key=os.environ["AFS_API_KEY"]) as afs: + workspace = await afs.workspace.create(name="foobar") + async with await afs.fs.mount( + workspaces=[{"name": workspace["name"]}], + mode="rw", + ) as fs: + await fs.write_file("/src/README.md", "hello world") + result = await fs.bash().exec("cat /foobar/src/README.md") + print(result.stdout) +``` + +Note that `afs.fs.mount(...)` is itself a coroutine, so the mounted filesystem +is entered with `async with await afs.fs.mount(...)`. + +The async client matters most inside an async server. A blocking HTTP call in +an `async` handler stalls the event loop and blocks every other request; the +async client `await`s its network I/O instead, so the loop stays free. + +```python +from fastapi import FastAPI +from redis_afs import AsyncAFS + +app = FastAPI() + +@app.get("/readme/{workspace}") +async def read_readme(workspace: str): + async with AsyncAFS() as afs: + async with await afs.fs.mount(workspaces=[{"name": workspace}], mode="ro") as fs: + return {"content": await fs.read_file(f"/{workspace}/README.md")} +``` + +Prefer `async with` (or call `await afs.aclose()` and `await fs.aclose()`) so +the underlying `httpx` connection pools are closed. + ## Authentication ```bash diff --git a/sdk/python/api-docs.md b/sdk/python/api-docs.md index efb5783..a0d5075 100644 --- a/sdk/python/api-docs.md +++ b/sdk/python/api-docs.md @@ -52,6 +52,71 @@ fs.bash().exec(command, ...) # run a shell command against the mount | Self-managed endpoint | `AFS_API_BASE_URL` or `AFS(base_url=...)` | | Available since | `redis-afs` `0.1.0` | +## Async API + +The async client mirrors the sync surface as coroutines, backed by an +`httpx.AsyncClient`. It requires the optional `async` extra +(`pip install 'redis-afs[async]'`). Importing `redis_afs` works without +`httpx`, but constructing an async client without it raises `AFSError` with an +install hint. + +### Syntax + +```python +from redis_afs import AsyncAFS + +afs = AsyncAFS( + api_key=None, + base_url=None, + timeout=30.0, + headers=None, +) +``` + +`AsyncAFS` takes the same parameters as `AFS`, with the same `AFS_API_KEY` and +`AFS_API_BASE_URL` environment fallbacks. + +### API Methods + +```python +await afs.workspace.create(...) # create a workspace +await afs.workspace.list() # list workspaces +await afs.workspace.get(workspace) # get one workspace +await afs.workspace.fork(...) # fork a workspace +await afs.workspace.delete(workspace) # delete a workspace + +await afs.checkpoint.list(workspace) # list checkpoints +await afs.checkpoint.create(...) # create a checkpoint +await afs.checkpoint.restore(...) # restore a checkpoint + +fs = await afs.fs.mount(...) # create an isolated SDK mount (async) +await fs.read_file(path) # read a text file +await fs.write_file(path, content) # write a text file +await fs.list_files(path, depth) # list a directory +await fs.glob(pattern, ...) # match paths +await fs.grep(pattern, **options) # search file contents +await fs.checkpoint(name) # checkpoint mounted workspaces +await fs.sync_from_remote() # materialize workspaces locally +await fs.sync_to_remote() # write modified local files back +await fs.bash().exec(command, ...) # run a shell command against the mount +``` + +`afs.fs.mount(...)` is a coroutine, so the mounted filesystem is entered with +`async with await afs.fs.mount(...)`. Both `AsyncAFS` and `AsyncMountedFS` are +async context managers and expose `await afs.aclose()` / `await fs.aclose()`; +prefer `async with` (or call `aclose()`) so the underlying `httpx` connection +pools are closed. `await afs.fs.mount(..., concurrency=16)` bounds the number of +in-flight network requests during file sync (default 16). + +### Async At A Glance + +| Field | Value | +| --- | --- | +| Import | `from redis_afs import AsyncAFS` | +| Requires | `redis-afs[async]` (installs `httpx`) | +| Transport | `httpx.AsyncClient` (MCP over HTTP) | +| Available since | `redis-afs` `0.1.0` | + ## Examples Create a workspace, write a file, run a command, and clean up the temporary @@ -461,3 +526,17 @@ fs.map_absolute_repo_paths(command) currently propagate local file deletion. - `MountedFS` does not yet expose `mkdir`, `rename`, `delete`, `stat`, streams, or binary file helpers. + +### Async Client + +- `AsyncMountedFS` bounds *in-flight network requests* during file sync with + the `concurrency` kwarg (default 16), but the number of pending coroutine + objects is not bounded, so memory scales with tree size. This is acceptable + for typical workspaces; a worker-pool design is a possible future + improvement. +- A failed `sync_from_remote()` (for example a mid-sync error) can leave a + partially-materialized local root, the same behavior as the sync client. +- Local disk I/O during sync stays synchronous; only network calls are async. + Using `aiofiles` is a possible future enhancement. +- `bash().exec(...)` raises `asyncio.TimeoutError` on timeout, whereas the sync + client raises `subprocess.TimeoutExpired`. diff --git a/sdk/python/pyproject.toml b/sdk/python/pyproject.toml index 226416b..c00c87f 100644 --- a/sdk/python/pyproject.toml +++ b/sdk/python/pyproject.toml @@ -32,6 +32,9 @@ classifiers = [ ] dependencies = [] +[project.optional-dependencies] +async = ["httpx>=0.27"] + [project.urls] Homepage = "https://github.com/redis/agent-filesystem/tree/main/sdk/python" Repository = "https://github.com/redis/agent-filesystem" diff --git a/sdk/python/src/redis_afs/__init__.py b/sdk/python/src/redis_afs/__init__.py index 2b92c66..f078b76 100644 --- a/sdk/python/src/redis_afs/__init__.py +++ b/sdk/python/src/redis_afs/__init__.py @@ -5,8 +5,19 @@ AFS, AFSError, MountedFS, + MountMode, WorkspaceClient, ) +from .aio import ( + AsyncAFS, + AsyncBashRunner, + AsyncCheckpointClient, + AsyncFSClient, + AsyncMCPHttpClient, + AsyncMountedFS, + AsyncRepoClient, + AsyncWorkspaceClient, +) __all__ = [ "BashResult", @@ -15,5 +26,14 @@ "AFS", "AFSError", "MountedFS", + "MountMode", "WorkspaceClient", + "AsyncAFS", + "AsyncBashRunner", + "AsyncCheckpointClient", + "AsyncFSClient", + "AsyncMCPHttpClient", + "AsyncMountedFS", + "AsyncRepoClient", + "AsyncWorkspaceClient", ] diff --git a/sdk/python/src/redis_afs/_mcp.py b/sdk/python/src/redis_afs/_mcp.py new file mode 100644 index 0000000..1cacda5 --- /dev/null +++ b/sdk/python/src/redis_afs/_mcp.py @@ -0,0 +1,47 @@ +"""MCP-over-HTTP wire helpers (endpoint, JSON-RPC framing, result unwrap).""" + +from __future__ import annotations + +import json +from typing import Any, Mapping + +from .errors import AFSError + + +def build_rpc_body(rpc_id: int, method: str, params: Mapping[str, Any] | None) -> bytes: + return json.dumps( + { + "jsonrpc": "2.0", + "id": rpc_id, + "method": method, + "params": dict(params or {}), + } + ).encode("utf-8") + + +def parse_rpc_payload(text: str) -> Any: + payload = json.loads(text or "{}") + if payload.get("error"): + error = payload["error"] + raise AFSError(str(error.get("message", "MCP request failed")), code=error.get("code"), payload=payload) + return payload.get("result") + + +def unwrap_tool_result(result: Any, name: str) -> Any: + if isinstance(result, dict) and result.get("isError"): + content = "\n".join(item.get("text", "") for item in result.get("content", [])) + raise AFSError(content or f"MCP tool {name} failed", payload=result) + if isinstance(result, dict): + return result.get("structuredContent", result) + return result + + +def normalize_mcp_endpoint(base_url: str) -> str: + trimmed = base_url.strip().rstrip("/") + if not trimmed: + raise AFSError("base_url is required") + return trimmed if trimmed.endswith("/mcp") else f"{trimmed}/mcp" + + +def strip_none(values: dict[str, Any]) -> dict[str, Any]: + return {key: value for key, value in values.items() if value is not None} diff --git a/sdk/python/src/redis_afs/_paths.py b/sdk/python/src/redis_afs/_paths.py new file mode 100644 index 0000000..5445e08 --- /dev/null +++ b/sdk/python/src/redis_afs/_paths.py @@ -0,0 +1,60 @@ +"""Remote path normalization and workspace path resolution.""" + +from __future__ import annotations + +import posixpath +import re +from pathlib import Path +from typing import Sequence + +from .errors import AFSError + + +def normalize_remote_path(path: str) -> str: + raw = path.strip() + if not raw: + return "/" + parts = [part for part in raw.split("/") if part] + if ".." in parts: + raise AFSError(f"path {path} must not contain '..'") + normalized = posixpath.normpath(raw if raw.startswith("/") else f"/{raw}") + return "/" if normalized == "." else normalized + + +class MountTable: + """Pure path logic for a set of mounted workspaces. + + Owns longest-prefix workspace resolution, the single-workspace fallback, + and the absolute-path rewrite used to map remote workspace roots onto a + local mirror. Holds no workspace clients and performs no I/O. + """ + + def __init__(self, workspace_names: Sequence[str]) -> None: + self._names = list(workspace_names) + + @property + def names(self) -> list[str]: + return list(self._names) + + def resolve(self, raw_path: str) -> tuple[str, str]: + """Resolve a raw path to ``(workspace_name, remote_path)``.""" + normalized = normalize_remote_path(raw_path) + for name in sorted(self._names, key=len, reverse=True): + prefix = f"/{name}" + if normalized == prefix: + return name, "/" + if normalized.startswith(f"{prefix}/"): + return name, normalized[len(prefix) :] or "/" + if len(self._names) == 1: + return self._names[0], normalized + choices = ", ".join(f"/{name}" for name in self._names) + raise AFSError(f"path {raw_path} must start with one of: {choices}") + + def map_absolute(self, command: str, local_root: str) -> str: + """Rewrite absolute ``/workspace`` prefixes to their local mirror paths.""" + out = command + for name in sorted(self._names, key=len, reverse=True): + remote_prefix = f"/{name}" + local_prefix = str(Path(local_root, name)).replace("\\", "/") + out = re.sub(rf"{re.escape(remote_prefix)}(?=/|\s|$)", local_prefix, out) + return out diff --git a/sdk/python/src/redis_afs/aio/__init__.py b/sdk/python/src/redis_afs/aio/__init__.py new file mode 100644 index 0000000..c9b9f69 --- /dev/null +++ b/sdk/python/src/redis_afs/aio/__init__.py @@ -0,0 +1,22 @@ +"""Async public surface for redis_afs, re-exported from the aio subpackage modules.""" + +from ._http import AsyncMCPHttpClient, httpx +from ._mount import AsyncBashRunner, AsyncMountedFS, _AsyncMountedWorkspace +from ._resources import ( + AsyncAFS, + AsyncCheckpointClient, + AsyncFSClient, + AsyncRepoClient, + AsyncWorkspaceClient, +) + +__all__ = [ + "AsyncMCPHttpClient", + "AsyncAFS", + "AsyncWorkspaceClient", + "AsyncRepoClient", + "AsyncCheckpointClient", + "AsyncFSClient", + "AsyncMountedFS", + "AsyncBashRunner", +] diff --git a/sdk/python/src/redis_afs/aio/_http.py b/sdk/python/src/redis_afs/aio/_http.py new file mode 100644 index 0000000..3ba8c15 --- /dev/null +++ b/sdk/python/src/redis_afs/aio/_http.py @@ -0,0 +1,83 @@ +"""Async HTTP/MCP transport: the httpx import guard and AsyncMCPHttpClient.""" + +from __future__ import annotations + +import os +from typing import Any, Mapping + +from ..errors import AFSError +from ..models import DEFAULT_BASE_URL +from .._mcp import ( + build_rpc_body, + normalize_mcp_endpoint, + parse_rpc_payload, + strip_none, + unwrap_tool_result, +) + +try: + import httpx +except ImportError: # pragma: no cover + httpx = None + +_INSTALL_HINT = "redis-afs[async] requires httpx; install with: pip install 'redis-afs[async]'" + + +class AsyncMCPHttpClient: + def __init__( + self, + *, + api_key: str | None = None, + base_url: str | None = None, + timeout: float = 30.0, + headers: Mapping[str, str] | None = None, + transport: Any | None = None, + ) -> None: + if httpx is None: + raise AFSError(_INSTALL_HINT) + self.api_key = api_key or os.environ.get("AFS_API_KEY") or "" + if not self.api_key: + raise AFSError("AFS api_key is required") + base = base_url or os.environ.get("AFS_API_BASE_URL") or DEFAULT_BASE_URL + self.endpoint = normalize_mcp_endpoint(base) + self.timeout = timeout + self.headers = dict(headers or {}) + self._next_id = 1 + self._client = httpx.AsyncClient(timeout=timeout, transport=transport) + + async def call_tool(self, name: str, arguments: Mapping[str, Any] | None = None) -> Any: + result = await self.request( + "tools/call", + {"name": name, "arguments": strip_none(dict(arguments or {}))}, + ) + return unwrap_tool_result(result, name) + + async def request(self, method: str, params: Mapping[str, Any] | None = None) -> Any: + body = build_rpc_body(self._next_id, method, params) + self._next_id += 1 + headers = { + "content-type": "application/json", + "authorization": f"Bearer {self.api_key}", + **self.headers, + } + try: + response = await self._client.post(self.endpoint, content=body, headers=headers) + except httpx.TimeoutException as exc: + raise AFSError(f"MCP request timed out after {self.timeout}s") from exc + if response.status_code >= 400: + text = response.text + raise AFSError( + f"MCP request failed with HTTP {response.status_code}: {text}", + status=response.status_code, + payload=text, + ) + return parse_rpc_payload(response.text) + + async def aclose(self) -> None: + await self._client.aclose() + + async def __aenter__(self) -> "AsyncMCPHttpClient": + return self + + async def __aexit__(self, exc_type: object, exc: object, tb: object) -> None: + await self.aclose() diff --git a/sdk/python/src/redis_afs/aio/_mount.py b/sdk/python/src/redis_afs/aio/_mount.py new file mode 100644 index 0000000..9f3a427 --- /dev/null +++ b/sdk/python/src/redis_afs/aio/_mount.py @@ -0,0 +1,219 @@ +"""Async mounted filesystem primitives: _AsyncMountedWorkspace, AsyncMountedFS, AsyncBashRunner.""" + +from __future__ import annotations + +import asyncio +import os +import shutil +import tempfile +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Mapping, MutableMapping, Sequence + +from ..errors import AFSError +from ..models import BashResult +from .._paths import MountTable, normalize_remote_path as _normalize_remote_path +from ._sync import _TreeSync + + +@dataclass(frozen=True) +class _AsyncMountedWorkspace: + name: str + token: str + client: Any + + +class AsyncMountedFS: + def __init__( + self, + workspaces: Sequence[_AsyncMountedWorkspace], + *, + mode: str = "rw", + concurrency: int = 16, + ) -> None: + self._workspaces = list(workspaces) + self._workspaces_by_name = {workspace.name: workspace for workspace in self._workspaces} + if len(self._workspaces_by_name) != len(self._workspaces): + raise AFSError("workspaces must be mounted at most once") + self.mode = mode + self._table = MountTable(self.workspace_names) + self._local_root: tempfile.TemporaryDirectory[str] | None = None + self._concurrency = concurrency + self._semaphore = asyncio.Semaphore(concurrency) + + @property + def repo_names(self) -> list[str]: + return self.workspace_names + + @property + def workspace_names(self) -> list[str]: + return [workspace.name for workspace in self._workspaces] + + @property + def local_root(self) -> str | None: + return self._local_root.name if self._local_root else None + + async def read_file(self, path: str) -> str: + workspace, remote_path = self._resolve(path) + response = await workspace.client.call_tool("file_read", {"path": remote_path}) + if response.get("binary"): + raise AFSError(f"file {remote_path} is binary and cannot be returned as text") + if response.get("kind") == "dir": + raise AFSError(f"path {remote_path} is a directory") + return str(response.get("content", "")) + + async def write_file(self, path: str, content: str | bytes) -> None: + workspace, remote_path = self._resolve(path) + text = content.decode("utf-8") if isinstance(content, bytes) else content + await workspace.client.call_tool("file_write", {"path": remote_path, "content": text}) + if self.local_root: + local_path = self._local_path_for(workspace.name, remote_path) + local_path.parent.mkdir(parents=True, exist_ok=True) + local_path.write_text(text, encoding="utf-8") + + async def list_files(self, path: str = "/", depth: int = 1) -> list[dict[str, Any]]: + workspace, remote_path = self._resolve(path) + response = await workspace.client.call_tool("file_list", {"path": remote_path, "depth": depth}) + return list(response.get("entries", [])) + + async def glob( + self, + pattern: str, + *, + path: str = "/", + kind: str | None = None, + limit: int | None = None, + ) -> dict[str, Any]: + workspace, remote_path = self._resolve(path) + return await workspace.client.call_tool( + "file_glob", + {"path": remote_path, "pattern": pattern, "kind": kind, "limit": limit}, + ) + + async def grep(self, pattern: str, **options: Any) -> dict[str, Any]: + workspace, remote_path = self._resolve(str(options.pop("path", "/"))) + return await workspace.client.call_tool("file_grep", {"path": remote_path, "pattern": pattern, **options}) + + async def checkpoint(self, name: str | None = None) -> list[dict[str, Any]]: + return [ + await workspace.client.call_tool("checkpoint_create", {"checkpoint": name}) + for workspace in self._workspaces + ] + + def bash(self) -> "AsyncBashRunner": + return AsyncBashRunner(self) + + async def sync_from_remote(self) -> str: + root = self._ensure_local_root() + for workspace in self._workspaces: + workspace_root = Path(root, workspace.name) + shutil.rmtree(workspace_root, ignore_errors=True) + workspace_root.mkdir(parents=True, exist_ok=True) + await _TreeSync(workspace.client, self._semaphore).pull("/", workspace_root) + return root + + async def sync_to_remote(self) -> None: + if not self.local_root: + return + for workspace in self._workspaces: + workspace_root = Path(self.local_root, workspace.name) + if workspace_root.exists(): + await _TreeSync(workspace.client, self._semaphore).push(workspace_root, "/") + + async def aclose(self) -> None: + results = await asyncio.gather( + *(workspace.client.aclose() for workspace in self._workspaces), + return_exceptions=True, + ) + if self._local_root: + self._local_root.cleanup() + self._local_root = None + for result in results: + if isinstance(result, BaseException): + raise result + + async def __aenter__(self) -> "AsyncMountedFS": + return self + + async def __aexit__(self, exc_type: object, exc: object, tb: object) -> None: + await self.aclose() + + def map_absolute_workspace_paths(self, command: str) -> str: + if not self.local_root: + return command + return self._table.map_absolute(command, self.local_root) + + def map_absolute_repo_paths(self, command: str) -> str: + return self.map_absolute_workspace_paths(command) + + def _resolve(self, raw_path: str) -> tuple[_AsyncMountedWorkspace, str]: + name, remote_path = self._table.resolve(raw_path) + return self._workspaces_by_name[name], remote_path + + def _ensure_local_root(self) -> str: + if not self._local_root: + self._local_root = tempfile.TemporaryDirectory(prefix="afs-fs-") + return self._local_root.name + + def _local_path_for(self, workspace_name: str, remote_path: str) -> Path: + if not self.local_root: + raise AFSError("mount has not been materialized locally yet") + relative = _normalize_remote_path(remote_path).lstrip("/") + return Path(self.local_root, workspace_name, relative) + + +class AsyncBashRunner: + def __init__(self, mounted_fs: AsyncMountedFS) -> None: + self._fs = mounted_fs + + async def exec( + self, + command: str, + *, + cwd: str | None = None, + env: Mapping[str, str | None] | None = None, + timeout: float | None = None, + check: bool = False, + ) -> BashResult: + # On timeout this raises asyncio.TimeoutError (an alias of builtin TimeoutError + # on Python 3.11+) after killing and reaping the spawned child. This is the + # async API's contract; the sync client raises subprocess.TimeoutExpired, but we + # intentionally surface the asyncio-native exception here. + root = await self._fs.sync_from_remote() + mapped_command = self._fs.map_absolute_workspace_paths(command) + run_env: MutableMapping[str, str] = dict(os.environ) + if env: + for key, value in env.items(): + if value is None: + run_env.pop(key, None) + else: + run_env[key] = value + proc = await asyncio.create_subprocess_exec( + "/bin/bash", + "-c", + mapped_command, + cwd=str(Path(root, cwd)) if cwd else root, + env=run_env, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + ) + if timeout is not None: + try: + stdout_b, stderr_b = await asyncio.wait_for(proc.communicate(), timeout) + except asyncio.TimeoutError: + proc.kill() + await proc.communicate() # reap the child and drain pipes + raise + else: + stdout_b, stderr_b = await proc.communicate() + await self._fs.sync_to_remote() + result = BashResult( + stdout=stdout_b.decode("utf-8"), + stderr=stderr_b.decode("utf-8"), + exit_code=proc.returncode, + command=command, + mapped_command=mapped_command, + ) + if check and result.exit_code != 0: + raise AFSError(f"command exited with status {result.exit_code}", payload=result) + return result diff --git a/sdk/python/src/redis_afs/aio/_resources.py b/sdk/python/src/redis_afs/aio/_resources.py new file mode 100644 index 0000000..a6e35f4 --- /dev/null +++ b/sdk/python/src/redis_afs/aio/_resources.py @@ -0,0 +1,202 @@ +"""Async resource clients: workspaces, checkpoints, filesystem mount, and the AsyncAFS facade.""" + +from __future__ import annotations + +import asyncio +from typing import Any, Mapping, Sequence + +from ..errors import AFSError +from ..models import ( + MountMode, + as_workspace_name as _workspace_name, +) +from ._http import AsyncMCPHttpClient +from ._mount import AsyncMountedFS, _AsyncMountedWorkspace + + +def _pick_workspace( + workspace: str | Mapping[str, Any] | None, + repo: str | Mapping[str, Any] | None, +) -> str: + """Resolve workspace-or-repo for workspace ops, coercing names (str|Mapping).""" + return _workspace_name(workspace if workspace is not None else repo) + + +def _require_checkpoint_workspace(workspace: str | None, repo: str | None, *, action: str) -> str: + """Resolve the checkpoint workspace by plain truthiness (no name coercion).""" + workspace_name = workspace or repo + if not workspace_name: + raise AFSError(f"checkpoint.{action} requires a workspace") + return workspace_name + + +class AsyncWorkspaceClient: + def __init__(self, mcp: "AsyncMCPHttpClient") -> None: + self._mcp = mcp + + async def create( + self, *, name: str, description: str | None = None, template_slug: str | None = None + ) -> dict[str, Any]: + return await self._mcp.call_tool( + "workspace_create", + { + "name": name, + "description": description, + "template_slug": template_slug, + }, + ) + + async def list(self) -> list[dict[str, Any]]: + response = await self._mcp.call_tool("workspace_list") + if isinstance(response, list): + return response + return list(response.get("items", [])) + + async def get( + self, + workspace: str | Mapping[str, Any] | None = None, + *, + repo: str | Mapping[str, Any] | None = None, + ) -> dict[str, Any]: + return await self._mcp.call_tool( + "workspace_get", + {"workspace": _pick_workspace(workspace, repo)}, + ) + + async def fork(self, *, source: str, name: str) -> dict[str, Any]: + return await self._mcp.call_tool("workspace_fork", {"source": source, "name": name}) + + async def delete( + self, + workspace: str | Mapping[str, Any] | None = None, + *, + repo: str | Mapping[str, Any] | None = None, + ) -> dict[str, Any]: + return await self._mcp.call_tool( + "workspace_delete", + {"workspace": _pick_workspace(workspace, repo)}, + ) + + +AsyncRepoClient = AsyncWorkspaceClient + + +class AsyncCheckpointClient: + def __init__(self, mcp: "AsyncMCPHttpClient") -> None: + self._mcp = mcp + + async def list(self, workspace: str | Mapping[str, Any]) -> list[dict[str, Any]]: + response = await self._mcp.call_tool("checkpoint_list", {"workspace": _workspace_name(workspace)}) + return list(response.get("checkpoints", [])) + + async def create( + self, + *, + workspace: str | None = None, + repo: str | None = None, + checkpoint: str | None = None, + ) -> dict[str, Any]: + workspace_name = _require_checkpoint_workspace(workspace, repo, action="create") + return await self._mcp.call_tool("checkpoint_create", {"workspace": workspace_name, "checkpoint": checkpoint}) + + async def restore( + self, *, workspace: str | None = None, repo: str | None = None, checkpoint: str + ) -> dict[str, Any]: + workspace_name = _require_checkpoint_workspace(workspace, repo, action="restore") + return await self._mcp.call_tool("checkpoint_restore", {"workspace": workspace_name, "checkpoint": checkpoint}) + + +class AsyncFSClient: + def __init__(self, control_plane: "AsyncMCPHttpClient") -> None: + self._control_plane = control_plane + + async def _mount_one( + self, + ref: Mapping[str, Any], + *, + profile: str, + token_name: str | None, + ) -> _AsyncMountedWorkspace: + name = _workspace_name(ref) + issued = await self._control_plane.call_tool( + "mcp_token_issue", + { + "workspace": name, + "name": token_name or f"redis-afs {name}", + "profile": profile, + }, + ) + token = str(issued.get("token", "")) + if not token: + raise AFSError(f"mcp_token_issue did not return a token for {name}", payload=issued) + return _AsyncMountedWorkspace( + name=name, + token=token, + client=AsyncMCPHttpClient( + api_key=token, + base_url=issued.get("url") or self._control_plane.endpoint, + timeout=self._control_plane.timeout, + ), + ) + + async def mount( + self, + *, + workspaces: Sequence[Mapping[str, Any]] | None = None, + repos: Sequence[Mapping[str, Any]] | None = None, + mode: MountMode | str = MountMode.RW, + token_name: str | None = None, + concurrency: int = 16, + ) -> "AsyncMountedFS": + workspace_refs = list(workspaces if workspaces is not None else repos or []) + if not workspace_refs: + raise AFSError("fs.mount requires at least one workspace") + profile = MountMode.coerce(mode).profile + # Issue every workspace token concurrently; gather preserves input order. + results = await asyncio.gather( + *(self._mount_one(ref, profile=profile, token_name=token_name) for ref in workspace_refs), + return_exceptions=True, + ) + mounted = [r for r in results if isinstance(r, _AsyncMountedWorkspace)] + failure = next((r for r in results if isinstance(r, BaseException)), None) + if failure is not None: + # Partial failure: close the children we did build before re-raising. + await asyncio.gather(*(m.client.aclose() for m in mounted), return_exceptions=True) + raise failure + return AsyncMountedFS(mounted, mode=mode, concurrency=concurrency) + + +class AsyncAFS: + def __init__( + self, + *, + api_key: str | None = None, + base_url: str | None = None, + timeout: float = 30.0, + headers: Mapping[str, str] | None = None, + ) -> None: + self._control_plane = AsyncMCPHttpClient( + api_key=api_key, + base_url=base_url, + timeout=timeout, + headers=headers, + ) + self.workspace = AsyncWorkspaceClient(self._control_plane) + self.workspaces = self.workspace + self.repo = self.workspace + self.repos = self.workspace + self.checkpoint = AsyncCheckpointClient(self._control_plane) + self.checkpoints = self.checkpoint + self.fs = AsyncFSClient(self._control_plane) + + async def call_tool(self, name: str, arguments: Mapping[str, Any] | None = None) -> Any: + return await self._control_plane.call_tool(name, arguments or {}) + + async def aclose(self) -> None: + await self._control_plane.aclose() + + async def __aenter__(self) -> "AsyncAFS": + return self + + async def __aexit__(self, exc_type: object, exc: object, tb: object) -> None: + await self.aclose() diff --git a/sdk/python/src/redis_afs/aio/_sync.py b/sdk/python/src/redis_afs/aio/_sync.py new file mode 100644 index 0000000..4a996b3 --- /dev/null +++ b/sdk/python/src/redis_afs/aio/_sync.py @@ -0,0 +1,87 @@ +"""Semaphore-bounded remote<->local tree copy for a single mounted workspace.""" + +from __future__ import annotations + +import asyncio +import posixpath +from pathlib import Path +from typing import Any, Awaitable, Callable + +from .._paths import normalize_remote_path as _normalize_remote_path + + +class _TreeSync: + """Copies one workspace's tree between the remote and a local directory. + + The semaphore guards only the network ``call_tool`` calls; local disk I/O + runs synchronously. Constructed per-workspace with a shared semaphore so the + concurrency bound spans every workspace in a mount. + """ + + def __init__(self, workspace_client: Any, semaphore: asyncio.Semaphore) -> None: + self._client = workspace_client + self._semaphore = semaphore + + async def pull(self, remote_root: str, local_dir: Path) -> None: + """Copy the remote tree at ``remote_root`` into ``local_dir`` (remote -> local).""" + await self._copy_remote_directory(remote_root, local_dir) + + async def push(self, local_dir: Path, remote_root: str) -> None: + """Copy ``local_dir`` into the remote tree at ``remote_root`` (local -> remote).""" + await self._copy_local_directory(local_dir, remote_root) + + async def _guarded(self, coro_factory: Callable[[], Awaitable[Any]]) -> Any: + # Takes a factory (callable returning a fresh coroutine), not a coroutine object, + # so the semaphore is acquired before the coroutine starts; do not "simplify" to a + # passed-in coroutine or it becomes a double-await/already-awaited bug. + async with self._semaphore: + return await coro_factory() + + async def _copy_remote_directory(self, remote_path: str, local_path: Path) -> None: + response = await self._guarded( + lambda: self._client.call_tool("file_list", {"path": remote_path, "depth": 1}) + ) + tasks: list[Awaitable[None]] = [] + for entry in response.get("entries", []): + target = local_path / entry["name"] + kind = entry.get("kind") + if kind == "dir": + # Create the directory before recursing so child writes land. + target.mkdir(parents=True, exist_ok=True) + tasks.append(self._copy_remote_directory(entry["path"], target)) + elif kind == "symlink" and entry.get("target"): + try: + target.symlink_to(entry["target"]) + except FileExistsError: + pass + elif kind == "file": + tasks.append(self._copy_remote_file(entry["path"], target)) + if tasks: + await asyncio.gather(*tasks) + + async def _copy_remote_file(self, remote_path: str, target: Path) -> None: + file_response = await self._guarded( + lambda: self._client.call_tool("file_read", {"path": remote_path}) + ) + if not file_response.get("binary"): + target.parent.mkdir(parents=True, exist_ok=True) + target.write_text(str(file_response.get("content", "")), encoding="utf-8") + + async def _copy_local_directory(self, local_directory: Path, remote_directory: str) -> None: + tasks: list[Awaitable[None]] = [] + for child in local_directory.iterdir(): + remote_path = _normalize_remote_path(posixpath.join(remote_directory, child.name)) + if child.is_symlink(): + continue + elif child.is_dir(): + tasks.append(self._copy_local_directory(child, remote_path)) + elif child.is_file(): + tasks.append(self._copy_local_file(child, remote_path)) + if tasks: + await asyncio.gather(*tasks) + + async def _copy_local_file(self, child: Path, remote_path: str) -> None: + content = child.read_text(encoding="utf-8") + await self._guarded( + lambda: self._client.call_tool("file_write", {"path": remote_path, "content": content}) + ) diff --git a/sdk/python/src/redis_afs/client.py b/sdk/python/src/redis_afs/client.py index 5382563..c96d4e5 100644 --- a/sdk/python/src/redis_afs/client.py +++ b/sdk/python/src/redis_afs/client.py @@ -1,6 +1,5 @@ from __future__ import annotations -import json import os import posixpath import re @@ -13,31 +12,21 @@ from pathlib import Path from typing import Any, Mapping, MutableMapping, Sequence -DEFAULT_BASE_URL = "https://afs.cloud" - - -class AFSError(RuntimeError): - def __init__( - self, - message: str, - *, - status: int | None = None, - code: int | None = None, - payload: Any | None = None, - ) -> None: - super().__init__(message) - self.status = status - self.code = code - self.payload = payload - - -@dataclass(frozen=True) -class BashResult: - stdout: str - stderr: str - exit_code: int - command: str - mapped_command: str +from .errors import AFSError +from .models import ( + DEFAULT_BASE_URL, + BashResult, + MountMode, + as_workspace_name as _workspace_name, +) +from ._mcp import ( + build_rpc_body, + normalize_mcp_endpoint as _normalize_mcp_endpoint, + parse_rpc_payload, + strip_none as _strip_none, + unwrap_tool_result, +) +from ._paths import normalize_remote_path as _normalize_remote_path class AFS: @@ -152,13 +141,13 @@ def mount( *, workspaces: Sequence[Mapping[str, Any]] | None = None, repos: Sequence[Mapping[str, Any]] | None = None, - mode: str = "rw", + mode: MountMode | str = MountMode.RW, token_name: str | None = None, ) -> "MountedFS": workspace_refs = list(workspaces if workspaces is not None else repos or []) if not workspace_refs: raise AFSError("fs.mount requires at least one workspace") - profile = _profile_for_mode(mode) + profile = MountMode.coerce(mode).profile mounted: list[_MountedWorkspace] = [] for workspace in workspace_refs: name = _workspace_name(workspace) @@ -432,20 +421,10 @@ def call_tool(self, name: str, arguments: Mapping[str, Any] | None = None) -> An "arguments": _strip_none(dict(arguments or {})), }, ) - if result.get("isError"): - content = "\n".join(item.get("text", "") for item in result.get("content", [])) - raise AFSError(content or f"MCP tool {name} failed", payload=result) - return result.get("structuredContent", result) + return unwrap_tool_result(result, name) def request(self, method: str, params: Mapping[str, Any] | None = None) -> Any: - body = json.dumps( - { - "jsonrpc": "2.0", - "id": self._next_id, - "method": method, - "params": dict(params or {}), - } - ).encode("utf-8") + body = build_rpc_body(self._next_id, method, params) self._next_id += 1 headers = { "content-type": "application/json", @@ -455,54 +434,8 @@ def request(self, method: str, params: Mapping[str, Any] | None = None) -> Any: request = urllib.request.Request(self.endpoint, data=body, headers=headers, method="POST") try: with urllib.request.urlopen(request, timeout=self.timeout) as response: - payload = json.loads(response.read().decode("utf-8") or "{}") + text = response.read().decode("utf-8") except urllib.error.HTTPError as exc: text = exc.read().decode("utf-8", errors="replace") raise AFSError(f"MCP request failed with HTTP {exc.code}: {text}", status=exc.code, payload=text) from exc - if payload.get("error"): - error = payload["error"] - raise AFSError(str(error.get("message", "MCP request failed")), code=error.get("code"), payload=payload) - return payload.get("result") - - -def _workspace_name(workspace: str | Mapping[str, Any] | None) -> str: - if isinstance(workspace, str): - return workspace - if workspace is None: - raise AFSError("workspace name is required") - name = str(workspace.get("name", "")).strip() - if not name: - raise AFSError("workspace name is required") - return name - - -def _profile_for_mode(mode: str) -> str: - if mode == "ro": - return "workspace-ro" - if mode == "rw": - return "workspace-rw" - if mode == "rw-checkpoint": - return "workspace-rw-checkpoint" - raise AFSError('mode must be "ro", "rw", or "rw-checkpoint"') - - -def _normalize_mcp_endpoint(base_url: str) -> str: - trimmed = base_url.strip().rstrip("/") - if not trimmed: - raise AFSError("base_url is required") - return trimmed if trimmed.endswith("/mcp") else f"{trimmed}/mcp" - - -def _normalize_remote_path(path: str) -> str: - raw = path.strip() - if not raw: - return "/" - parts = [part for part in raw.split("/") if part] - if ".." in parts: - raise AFSError(f"path {path} must not contain '..'") - normalized = posixpath.normpath(raw if raw.startswith("/") else f"/{raw}") - return "/" if normalized == "." else normalized - - -def _strip_none(values: dict[str, Any]) -> dict[str, Any]: - return {key: value for key, value in values.items() if value is not None} + return parse_rpc_payload(text) diff --git a/sdk/python/src/redis_afs/errors.py b/sdk/python/src/redis_afs/errors.py new file mode 100644 index 0000000..589cdc6 --- /dev/null +++ b/sdk/python/src/redis_afs/errors.py @@ -0,0 +1,20 @@ +"""Public exception types for the redis_afs SDK.""" + +from __future__ import annotations + +from typing import Any + + +class AFSError(RuntimeError): + def __init__( + self, + message: str, + *, + status: int | None = None, + code: int | None = None, + payload: Any | None = None, + ) -> None: + super().__init__(message) + self.status = status + self.code = code + self.payload = payload diff --git a/sdk/python/src/redis_afs/models.py b/sdk/python/src/redis_afs/models.py new file mode 100644 index 0000000..25498a2 --- /dev/null +++ b/sdk/python/src/redis_afs/models.py @@ -0,0 +1,55 @@ +"""Public data models and shared constants for the redis_afs SDK.""" + +from __future__ import annotations + +from dataclasses import dataclass +from enum import Enum +from typing import Any, Mapping + +from .errors import AFSError + +DEFAULT_BASE_URL = "https://afs.cloud" + + +class MountMode(str, Enum): + """Workspace mount mode. + + Mirrors the control plane's workspace-* MCP profiles + (MCPProfileWorkspaceRO/RW/RWCheckpoint); the profile is always + ``workspace-{mode}``. + """ + + RO = "ro" + RW = "rw" + RW_CHECKPOINT = "rw-checkpoint" + + @property + def profile(self) -> str: + return f"workspace-{self.value}" + + @classmethod + def coerce(cls, value: "MountMode | str") -> "MountMode": + try: + return cls(value) + except ValueError: + raise AFSError(f"mode must be one of {[m.value for m in cls]}, got {value!r}") + + +@dataclass(frozen=True) +class BashResult: + stdout: str + stderr: str + exit_code: int + command: str + mapped_command: str + + +def as_workspace_name(workspace: str | Mapping[str, Any] | None) -> str: + if isinstance(workspace, str): + return workspace + if workspace is None: + raise AFSError("workspace name is required") + name = str(workspace.get("name", "")).strip() + if not name: + raise AFSError("workspace name is required") + return name diff --git a/sdk/python/tests/test_aio.py b/sdk/python/tests/test_aio.py new file mode 100644 index 0000000..7b920da --- /dev/null +++ b/sdk/python/tests/test_aio.py @@ -0,0 +1,466 @@ +import asyncio +import unittest +from pathlib import Path +from unittest.mock import patch + +from redis_afs.errors import AFSError +from redis_afs.models import BashResult + +try: + import httpx +except ImportError: # pragma: no cover + httpx = None + +import redis_afs.aio as aio +from redis_afs.aio import ( + AsyncAFS, + AsyncCheckpointClient, + AsyncFSClient, + AsyncMCPHttpClient, + AsyncMountedFS, + AsyncRepoClient, + AsyncWorkspaceClient, + _AsyncMountedWorkspace, +) + + +class FakeAsyncMCP: + def __init__(self): + self.files = {} + self.symlinks = {} + self.calls = [] + self.issued = [] + + async def call_tool(self, name, arguments=None): + arguments = arguments or {} + self.calls.append((name, dict(arguments))) + if name == "mcp_token_issue": + self.issued.append(dict(arguments)) + return { + "token": f"token-{arguments['workspace']}", + "url": "https://afs.example/mcp", + "workspace": arguments["workspace"], + "profile": arguments["profile"], + } + if name == "workspace_list": + return {"items": [{"name": "a"}, {"name": "b"}]} + if name == "workspace_get": + return {"name": arguments["workspace"]} + if name == "workspace_delete": + return {"workspace": arguments["workspace"], "deleted": True} + if name == "checkpoint_list": + return {"checkpoints": [{"id": "1", "name": "head"}]} + if name == "checkpoint_create": + return {"workspace": arguments.get("workspace"), "checkpoint": arguments.get("checkpoint") or "auto", "created": True} + if name == "checkpoint_restore": + return {"workspace": arguments.get("workspace"), "checkpoint": arguments["checkpoint"], "restored": True} + # file ops reused by later tasks + if name == "file_write": + self.files[arguments["path"]] = arguments["content"] + return {"path": arguments["path"], "operation": "write"} + if name == "file_read": + if arguments["path"] in self.symlinks: + return {"path": arguments["path"], "kind": "symlink", "target": self.symlinks[arguments["path"]]} + return {"path": arguments["path"], "kind": "file", "content": self.files.get(arguments["path"], "")} + if name == "file_list": + return {"entries": _fake_entries(self.files, self.symlinks, arguments.get("path", "/"))} + raise AssertionError(f"unexpected tool {name}") + + async def aclose(self): + pass + + +def _fake_entries(files, symlinks, path): + entries = [] + for p in sorted(files): + if path == "/" and "/" not in p.strip("/"): + entries.append({"path": p, "name": p.strip("/"), "kind": "file"}) + elif p.startswith(path.rstrip("/") + "/"): + rem = p[len(path.rstrip("/")) + 1:] + if "/" not in rem: + entries.append({"path": p, "name": rem, "kind": "file"}) + for lp, target in sorted(symlinks.items()): + if path == "/" and "/" not in lp.strip("/"): + entries.append({"path": lp, "name": lp.strip("/"), "kind": "symlink", "target": target}) + return entries + + +class AsyncClientsTest(unittest.IsolatedAsyncioTestCase): + async def test_workspace_list_normalizes_items(self): + ws = AsyncWorkspaceClient(FakeAsyncMCP()) + self.assertEqual(await ws.list(), [{"name": "a"}, {"name": "b"}]) + + async def test_checkpoint_round_trip(self): + cp = AsyncCheckpointClient(FakeAsyncMCP()) + created = await cp.create(workspace="repo", checkpoint="c1") + restored = await cp.restore(workspace="repo", checkpoint="c1") + self.assertTrue(created["created"]) + self.assertTrue(restored["restored"]) + + async def test_checkpoint_create_requires_workspace(self): + cp = AsyncCheckpointClient(FakeAsyncMCP()) + with self.assertRaises(AFSError): + await cp.create() + + async def test_checkpoint_restore_requires_workspace(self): + cp = AsyncCheckpointClient(FakeAsyncMCP()) + with self.assertRaises(AFSError): + await cp.restore(checkpoint="c1") + + async def test_workspace_get_resolves_repo_argument(self): + ws = AsyncWorkspaceClient(FakeAsyncMCP()) + # AsyncRepoClient is an alias of AsyncWorkspaceClient. + self.assertIs(AsyncRepoClient, AsyncWorkspaceClient) + self.assertEqual(await ws.get(repo="myrepo"), {"name": "myrepo"}) + + async def test_workspace_delete_resolves_workspace_argument(self): + fake = FakeAsyncMCP() + ws = AsyncWorkspaceClient(fake) + result = await ws.delete("myws") + self.assertEqual(result, {"workspace": "myws", "deleted": True}) + + +class ExportsTest(unittest.TestCase): + def test_async_names_importable_from_package_root(self): + import redis_afs + for name in ["AsyncAFS", "AsyncWorkspaceClient", "AsyncRepoClient", + "AsyncCheckpointClient", "AsyncFSClient", "AsyncMountedFS", + "AsyncBashRunner", "AsyncMCPHttpClient"]: + self.assertTrue(hasattr(redis_afs, name), name) + + +class HttpxGuardTest(unittest.TestCase): + def test_constructing_client_without_httpx_raises_install_hint(self): + original = aio._http.httpx + try: + aio._http.httpx = None + with self.assertRaises(AFSError) as ctx: + AsyncMCPHttpClient(api_key="k") + finally: + aio._http.httpx = original + self.assertIn("httpx", str(ctx.exception)) + + +def _mock_client(handler): + transport = httpx.MockTransport(handler) + return AsyncMCPHttpClient(api_key="k", base_url="https://afs.example", transport=transport) + + +@unittest.skipIf(httpx is None, "httpx not installed") +class AsyncTransportTest(unittest.IsolatedAsyncioTestCase): + async def test_call_tool_unwraps_structured_content(self): + seen = {} + + def handler(request): + seen["url"] = str(request.url) + seen["body"] = request.content.decode() + return httpx.Response(200, json={"result": {"structuredContent": {"ok": True}}}) + + client = _mock_client(handler) + try: + result = await client.call_tool("workspace_list", {"a": None, "b": 1}) + finally: + await client.aclose() + + self.assertEqual(result, {"ok": True}) + self.assertTrue(seen["url"].endswith("/mcp")) + self.assertNotIn('"a"', seen["body"]) # strip_none drops a=None + + async def test_tool_error_raises(self): + def handler(request): + return httpx.Response(200, json={"result": {"isError": True, "content": [{"text": "boom"}]}}) + + client = _mock_client(handler) + with self.assertRaises(AFSError) as ctx: + await client.call_tool("x") + await client.aclose() + self.assertIn("boom", str(ctx.exception)) + + async def test_http_error_carries_status(self): + def handler(request): + return httpx.Response(500, text="nope") + + client = _mock_client(handler) + with self.assertRaises(AFSError) as ctx: + await client.request("tools/call", {}) + await client.aclose() + self.assertEqual(ctx.exception.status, 500) + + async def test_timeout_maps_to_afs_error(self): + def handler(request): + raise httpx.TimeoutException("slow") + + client = _mock_client(handler) + with self.assertRaises(AFSError) as ctx: + await client.request("tools/call", {}) + await client.aclose() + self.assertIn("timed out", str(ctx.exception)) + + async def test_async_context_manager_closes(self): + def handler(request): + return httpx.Response(200, json={"result": {"structuredContent": {}}}) + + async with _mock_client(handler) as client: + await client.call_tool("x") + self.assertTrue(client._client.is_closed) + + +class AsyncMountedFSTest(unittest.IsolatedAsyncioTestCase): + async def test_single_workspace_paths_are_workspace_relative(self): + fake = FakeAsyncMCP() + fs = AsyncMountedFS([_AsyncMountedWorkspace(name="foobar", token="t", client=fake)]) + await fs.write_file("/src/README.md", "hello") + self.assertEqual(fake.files["/src/README.md"], "hello") + self.assertEqual(await fs.read_file("/foobar/src/README.md"), "hello") + self.assertEqual(fs.workspace_names, ["foobar"]) + + async def test_multi_workspace_requires_prefix(self): + fs = AsyncMountedFS([ + _AsyncMountedWorkspace(name="api", token="t", client=FakeAsyncMCP()), + _AsyncMountedWorkspace(name="web", token="t", client=FakeAsyncMCP()), + ]) + with self.assertRaises(AFSError): + await fs.write_file("/README.md", "hello") + + +class MountTableTest(unittest.TestCase): + def test_single_workspace_fallback(self): + from redis_afs._paths import MountTable + table = MountTable(["only"]) + self.assertEqual(table.resolve("/src/app.py"), ("only", "/src/app.py")) + + def test_exact_prefix_match(self): + from redis_afs._paths import MountTable + table = MountTable(["api", "web"]) + self.assertEqual(table.resolve("/api"), ("api", "/")) + self.assertEqual(table.resolve("/web/index.html"), ("web", "/index.html")) + + def test_multi_workspace_requires_prefix(self): + from redis_afs._paths import MountTable + table = MountTable(["api", "web"]) + with self.assertRaises(AFSError) as ctx: + table.resolve("/README.md") + self.assertIn("must start with one of", str(ctx.exception)) + self.assertIn("/api", str(ctx.exception)) + self.assertIn("/web", str(ctx.exception)) + + +class FakeAsyncControlPlane: + def __init__(self): + self.issued = [] + self.timeout = 30.0 + self.endpoint = "https://afs.example/mcp" + + async def call_tool(self, name, arguments=None): + arguments = arguments or {} + if name != "mcp_token_issue": + raise AssertionError(f"unexpected tool {name}") + token = f"workspace-token-{arguments['workspace']}" + self.issued.append({"name": name, "arguments": dict(arguments), "token": token}) + return { + "token": token, + "url": "https://afs.example/mcp", + "workspace": arguments["workspace"], + "profile": arguments["profile"], + } + + +class FakeAsyncMountedClient: + files_by_token = {} + + def __init__(self, *, api_key, base_url=None, timeout=30.0, headers=None): + self.api_key = api_key + self.endpoint = base_url or "https://afs.example/mcp" + self.timeout = timeout + self.headers = dict(headers or {}) + self.closed = False + + async def call_tool(self, name, arguments=None): + arguments = arguments or {} + files = self.files_by_token.setdefault(self.api_key, {}) + if name == "file_write": + files[arguments["path"]] = arguments["content"] + return {"path": arguments["path"], "operation": "write"} + if name == "file_read": + return {"path": arguments["path"], "kind": "file", "content": files.get(arguments["path"], "")} + raise AssertionError(f"unexpected tool {name}") + + async def aclose(self): + self.closed = True + + +class AsyncFSMountTest(unittest.IsolatedAsyncioTestCase): + async def test_mount_issues_token_and_round_trips(self): + control_plane = FakeAsyncControlPlane() + FakeAsyncMountedClient.files_by_token = {} + + with patch("redis_afs.aio._resources.AsyncMCPHttpClient", FakeAsyncMountedClient): + fs = await AsyncFSClient(control_plane).mount( + workspaces=[{"name": "repo"}], mode="rw", token_name="Mounted FS" + ) + try: + await fs.write_file("/repo/README.md", "hello from mounted fs") + self.assertEqual(await fs.read_file("/repo/README.md"), "hello from mounted fs") + self.assertEqual(fs.workspace_names, ["repo"]) + self.assertEqual(control_plane.issued[0]["arguments"]["workspace"], "repo") + self.assertEqual(control_plane.issued[0]["arguments"]["profile"], "workspace-rw") + self.assertEqual(control_plane.issued[0]["arguments"]["name"], "Mounted FS") + child = fs._workspaces[0].client + self.assertFalse(child.closed) + finally: + await fs.aclose() + self.assertTrue(child.closed) + + async def test_mount_issues_distinct_tokens_per_workspace(self): + control_plane = FakeAsyncControlPlane() + FakeAsyncMountedClient.files_by_token = {} + + with patch("redis_afs.aio._resources.AsyncMCPHttpClient", FakeAsyncMountedClient): + fs = await AsyncFSClient(control_plane).mount( + workspaces=[{"name": "api"}, {"name": "web"}], mode="rw" + ) + try: + self.assertEqual(fs.workspace_names, ["api", "web"]) + self.assertEqual(len(control_plane.issued), 2) + self.assertEqual(control_plane.issued[0]["arguments"]["workspace"], "api") + self.assertEqual(control_plane.issued[1]["arguments"]["workspace"], "web") + finally: + await fs.aclose() + + async def test_mount_issues_all_tokens_concurrently(self): + control_plane = FakeAsyncControlPlane() + FakeAsyncMountedClient.files_by_token = {} + + with patch("redis_afs.aio._resources.AsyncMCPHttpClient", FakeAsyncMountedClient): + fs = await AsyncFSClient(control_plane).mount( + workspaces=[{"name": "api"}, {"name": "web"}, {"name": "db"}], mode="rw" + ) + try: + self.assertEqual(len(control_plane.issued), 3) + # mounted ordering is deterministic, matching the input refs. + self.assertEqual(fs.workspace_names, ["api", "web", "db"]) + finally: + await fs.aclose() + + async def test_mount_requires_at_least_one_workspace(self): + with self.assertRaises(AFSError): + await AsyncFSClient(FakeAsyncControlPlane()).mount(workspaces=[]) + + async def test_mount_forwards_concurrency(self): + control_plane = FakeAsyncControlPlane() + FakeAsyncMountedClient.files_by_token = {} + + with patch("redis_afs.aio._resources.AsyncMCPHttpClient", FakeAsyncMountedClient): + fs = await AsyncFSClient(control_plane).mount( + workspaces=[{"name": "repo"}], concurrency=3 + ) + try: + self.assertEqual(fs._concurrency, 3) + finally: + await fs.aclose() + + +class AsyncSyncTest(unittest.IsolatedAsyncioTestCase): + async def test_round_trip_materializes_files(self): + fake = FakeAsyncMCP() + fake.files["/README.md"] = "hello" + fs = AsyncMountedFS([_AsyncMountedWorkspace(name="repo", token="t", client=fake)]) + self.addAsyncCleanup(fs.aclose) + root = await fs.sync_from_remote() + self.assertEqual(Path(root, "repo", "README.md").read_text(), "hello") + + async def test_sync_to_remote_skips_symlinks(self): + fake = FakeAsyncMCP() + fake.files["/README.md"] = "hello" + fake.symlinks["/link.md"] = "README.md" + fs = AsyncMountedFS([_AsyncMountedWorkspace(name="repo", token="t", client=fake)]) + self.addAsyncCleanup(fs.aclose) + root = await fs.sync_from_remote() + self.assertTrue(Path(root, "repo", "link.md").is_symlink()) + await fs.sync_to_remote() + writes = [a["path"] for n, a in fake.calls if n == "file_write"] + self.assertNotIn("/link.md", writes) + + +class ConcurrencyFakeMCP(FakeAsyncMCP): + def __init__(self): + super().__init__() + self.in_flight = 0 + self.max_in_flight = 0 + + async def call_tool(self, name, arguments=None): + if name == "file_read": + self.in_flight += 1 + self.max_in_flight = max(self.max_in_flight, self.in_flight) + try: + # Yield repeatedly so other concurrent reads can enter the + # critical section before we decrement, exposing real overlap. + for _ in range(5): + await asyncio.sleep(0) + return await super().call_tool(name, arguments) + finally: + self.in_flight -= 1 + return await super().call_tool(name, arguments) + + +class AsyncSyncConcurrencyTest(unittest.IsolatedAsyncioTestCase): + async def test_sync_from_remote_bounds_and_overlaps_reads(self): + fake = ConcurrencyFakeMCP() + for i in range(8): + fake.files[f"/file{i}.txt"] = f"content-{i}" + fs = AsyncMountedFS( + [_AsyncMountedWorkspace(name="repo", token="t", client=fake)], + concurrency=3, + ) + self.addAsyncCleanup(fs.aclose) + await fs.sync_from_remote() + self.assertLessEqual(fake.max_in_flight, 3) + self.assertGreater(fake.max_in_flight, 1) + + +class AsyncBashTest(unittest.IsolatedAsyncioTestCase): + async def test_exec_runs_command_and_syncs(self): + fake = FakeAsyncMCP() + fake.files["/hello.txt"] = "hi" + fs = AsyncMountedFS([_AsyncMountedWorkspace(name="repo", token="t", client=fake)]) + self.addAsyncCleanup(fs.aclose) + result = await fs.bash().exec("cat /repo/hello.txt") + self.assertIsInstance(result, BashResult) + self.assertEqual(result.exit_code, 0) + self.assertEqual(result.stdout.strip(), "hi") + + async def test_exec_check_raises_on_nonzero(self): + fake = FakeAsyncMCP() + fs = AsyncMountedFS([_AsyncMountedWorkspace(name="repo", token="t", client=fake)]) + self.addAsyncCleanup(fs.aclose) + # Without check=True, a non-zero exit is returned, not raised. + result = await fs.bash().exec("exit 3") + self.assertIsInstance(result, BashResult) + self.assertEqual(result.exit_code, 3) + # With check=True, it raises AFSError carrying the result as payload. + with self.assertRaises(AFSError) as ctx: + await fs.bash().exec("exit 3", check=True) + self.assertIs(ctx.exception.payload.exit_code, 3) + + async def test_exec_env_override(self): + fake = FakeAsyncMCP() + fs = AsyncMountedFS([_AsyncMountedWorkspace(name="repo", token="t", client=fake)]) + self.addAsyncCleanup(fs.aclose) + result = await fs.bash().exec("echo $FOO", env={"FOO": "bar"}) + self.assertEqual(result.exit_code, 0) + self.assertIn("bar", result.stdout) + # A value of None removes the variable from the environment. + removed = await fs.bash().exec("echo [$FOO]", env={"FOO": None}) + self.assertEqual(removed.stdout.strip(), "[]") + + async def test_exec_timeout_raises_and_reaps(self): + import asyncio + fake = FakeAsyncMCP() + fs = AsyncMountedFS([_AsyncMountedWorkspace(name="repo", token="t", client=fake)]) + self.addAsyncCleanup(fs.aclose) + with self.assertRaises(asyncio.TimeoutError): + await fs.bash().exec("sleep 5", timeout=0.2) + + +if __name__ == "__main__": + unittest.main()