initial a2a-pack
This commit is contained in:
446
a2a_pack/workspace.py
Normal file
446
a2a_pack/workspace.py
Normal file
@@ -0,0 +1,446 @@
|
||||
"""Workspace capability negotiation.
|
||||
|
||||
Agents never receive a filesystem path. They negotiate a *view* by intent::
|
||||
|
||||
view = await ctx.workspace.open_view(
|
||||
purpose="Fix failing payment test",
|
||||
hints=["payment", "checkout"],
|
||||
file_types=["python"],
|
||||
max_files=10,
|
||||
mode=WorkspaceMode.READ_WRITE_OVERLAY,
|
||||
)
|
||||
for path in view.files:
|
||||
content = await view.read(path)
|
||||
|
||||
The runtime resolves the request (semantic search + dependency graph + git
|
||||
metadata + policy + optional human approval) and returns a bounded grant.
|
||||
Writes are staged as :class:`WorkspacePatch` objects, never applied directly
|
||||
to the host filesystem from inside the sandbox.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
from abc import ABC, abstractmethod
|
||||
from datetime import datetime
|
||||
from enum import Enum
|
||||
from typing import Literal, Sequence
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, NonNegativeInt, PositiveInt
|
||||
|
||||
|
||||
class WorkspaceMode(str, Enum):
|
||||
READ_ONLY = "read_only"
|
||||
READ_WRITE_OVERLAY = "read_write_overlay" # writes staged as patches
|
||||
READ_WRITE_DIRECT = "read_write_direct" # discouraged; needs explicit policy
|
||||
|
||||
|
||||
class FileType(str, Enum):
|
||||
PYTHON = "python"
|
||||
TYPESCRIPT = "typescript"
|
||||
JAVASCRIPT = "javascript"
|
||||
YAML = "yaml"
|
||||
JSON = "json"
|
||||
TOML = "toml"
|
||||
MARKDOWN = "markdown"
|
||||
SQL = "sql"
|
||||
SHELL = "shell"
|
||||
OTHER = "other"
|
||||
|
||||
|
||||
class FileMatch(BaseModel):
|
||||
"""Result row from :meth:`WorkspaceClient.search`."""
|
||||
|
||||
model_config = ConfigDict(extra="forbid", frozen=True)
|
||||
|
||||
path: str
|
||||
file_type: FileType
|
||||
score: float = 0.0
|
||||
summary: str | None = None
|
||||
size_bytes: NonNegativeInt = 0
|
||||
|
||||
|
||||
class WorkspaceGrant(BaseModel):
|
||||
"""An approved access grant for a bounded set of files."""
|
||||
|
||||
model_config = ConfigDict(extra="forbid", frozen=True)
|
||||
|
||||
grant_id: str
|
||||
purpose: str
|
||||
files: tuple[FileMatch, ...]
|
||||
mode: WorkspaceMode
|
||||
reason: str
|
||||
expires_at: datetime | None = None
|
||||
requires_human_approval: bool = False
|
||||
|
||||
|
||||
class WorkspacePatch(BaseModel):
|
||||
"""A staged write. Not applied until the runtime/approver commits it."""
|
||||
|
||||
model_config = ConfigDict(extra="forbid", frozen=True)
|
||||
|
||||
grant_id: str
|
||||
path: str
|
||||
operation: Literal["create", "update", "delete"]
|
||||
content: bytes | None = None # bytes for create/update, None for delete
|
||||
|
||||
|
||||
class WorkspaceAccess(BaseModel):
|
||||
"""Class-level workspace policy.
|
||||
|
||||
Use :meth:`none` for agents that do not touch any workspace, or
|
||||
:meth:`dynamic` to allow capability negotiation under bounds.
|
||||
"""
|
||||
|
||||
model_config = ConfigDict(extra="forbid", frozen=True)
|
||||
|
||||
enabled: bool = False
|
||||
max_files: NonNegativeInt = 0
|
||||
allowed_modes: tuple[WorkspaceMode, ...] = ()
|
||||
require_reason: bool = True
|
||||
deny_patterns: tuple[str, ...] = ()
|
||||
require_human_approval: bool = False
|
||||
max_total_size_bytes: PositiveInt = 100 * 1024 * 1024
|
||||
|
||||
@classmethod
|
||||
def none(cls) -> "WorkspaceAccess":
|
||||
return cls(enabled=False)
|
||||
|
||||
@classmethod
|
||||
def dynamic(
|
||||
cls,
|
||||
*,
|
||||
max_files: int = 25,
|
||||
allowed_modes: Sequence[WorkspaceMode] = (WorkspaceMode.READ_ONLY,),
|
||||
require_reason: bool = True,
|
||||
deny_patterns: Sequence[str] = (),
|
||||
require_human_approval: bool = False,
|
||||
max_total_size_bytes: int = 100 * 1024 * 1024,
|
||||
) -> "WorkspaceAccess":
|
||||
return cls(
|
||||
enabled=True,
|
||||
max_files=max_files,
|
||||
allowed_modes=tuple(allowed_modes),
|
||||
require_reason=require_reason,
|
||||
deny_patterns=tuple(deny_patterns),
|
||||
require_human_approval=require_human_approval,
|
||||
max_total_size_bytes=max_total_size_bytes,
|
||||
)
|
||||
|
||||
|
||||
class WorkspaceDenied(PermissionError):
|
||||
"""Raised when a workspace request violates the agent's policy."""
|
||||
|
||||
|
||||
class WorkspaceView(ABC):
|
||||
"""A bounded view over a granted set of files.
|
||||
|
||||
Returned by :meth:`WorkspaceClient.open_view`. Reads always go to the
|
||||
granted view; writes return :class:`WorkspacePatch` objects that the
|
||||
runtime will commit (or reject) outside the sandbox.
|
||||
"""
|
||||
|
||||
grant: WorkspaceGrant
|
||||
|
||||
@property
|
||||
def files(self) -> tuple[FileMatch, ...]:
|
||||
return self.grant.files
|
||||
|
||||
@abstractmethod
|
||||
async def read(self, path: str) -> bytes: ...
|
||||
|
||||
@abstractmethod
|
||||
async def write(self, path: str, content: bytes) -> WorkspacePatch: ...
|
||||
|
||||
@abstractmethod
|
||||
async def delete(self, path: str) -> WorkspacePatch: ...
|
||||
|
||||
@abstractmethod
|
||||
async def patches(self) -> tuple[WorkspacePatch, ...]: ...
|
||||
|
||||
|
||||
class WorkspaceClient(ABC):
|
||||
"""Negotiation surface handed to the agent via ``ctx.workspace``.
|
||||
|
||||
The concrete implementation is provided by the runtime; agents must
|
||||
program against this interface only.
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
async def search(
|
||||
self,
|
||||
*,
|
||||
query: str,
|
||||
types: Sequence[FileType] = (),
|
||||
limit: int = 20,
|
||||
) -> list[FileMatch]: ...
|
||||
|
||||
@abstractmethod
|
||||
async def request_access(
|
||||
self,
|
||||
*,
|
||||
files: Sequence[FileMatch | str],
|
||||
mode: WorkspaceMode,
|
||||
reason: str,
|
||||
purpose: str = "",
|
||||
) -> WorkspaceGrant: ...
|
||||
|
||||
@abstractmethod
|
||||
async def open_view(
|
||||
self,
|
||||
*,
|
||||
purpose: str,
|
||||
hints: Sequence[str] = (),
|
||||
file_types: Sequence[FileType] = (),
|
||||
max_files: int = 10,
|
||||
mode: WorkspaceMode = WorkspaceMode.READ_ONLY,
|
||||
reason: str | None = None,
|
||||
) -> WorkspaceView: ...
|
||||
|
||||
@abstractmethod
|
||||
async def list_grants(self) -> list[WorkspaceGrant]: ...
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Local in-memory implementation, for dev/tests.
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class LocalWorkspaceView(WorkspaceView):
|
||||
def __init__(
|
||||
self,
|
||||
grant: WorkspaceGrant,
|
||||
client: "LocalWorkspaceClient",
|
||||
) -> None:
|
||||
self.grant = grant
|
||||
self._client = client
|
||||
self._patches: list[WorkspacePatch] = []
|
||||
self._granted_paths = {f.path for f in grant.files}
|
||||
|
||||
def _check(self, path: str) -> None:
|
||||
if path not in self._granted_paths:
|
||||
raise WorkspaceDenied(
|
||||
f"path {path!r} not in grant {self.grant.grant_id}"
|
||||
)
|
||||
|
||||
async def read(self, path: str) -> bytes:
|
||||
self._check(path)
|
||||
return self._client._files[path]
|
||||
|
||||
async def write(self, path: str, content: bytes) -> WorkspacePatch:
|
||||
if self.grant.mode is WorkspaceMode.READ_ONLY:
|
||||
raise WorkspaceDenied(f"grant is read-only: {self.grant.grant_id}")
|
||||
self._check(path)
|
||||
op: Literal["create", "update"] = (
|
||||
"create" if path not in self._client._files else "update"
|
||||
)
|
||||
patch = WorkspacePatch(
|
||||
grant_id=self.grant.grant_id,
|
||||
path=path,
|
||||
operation=op,
|
||||
content=content,
|
||||
)
|
||||
self._patches.append(patch)
|
||||
return patch
|
||||
|
||||
async def delete(self, path: str) -> WorkspacePatch:
|
||||
if self.grant.mode is WorkspaceMode.READ_ONLY:
|
||||
raise WorkspaceDenied(f"grant is read-only: {self.grant.grant_id}")
|
||||
self._check(path)
|
||||
patch = WorkspacePatch(
|
||||
grant_id=self.grant.grant_id,
|
||||
path=path,
|
||||
operation="delete",
|
||||
content=None,
|
||||
)
|
||||
self._patches.append(patch)
|
||||
return patch
|
||||
|
||||
async def patches(self) -> tuple[WorkspacePatch, ...]:
|
||||
return tuple(self._patches)
|
||||
|
||||
|
||||
class LocalWorkspaceClient(WorkspaceClient):
|
||||
"""In-memory workspace for local dev and tests.
|
||||
|
||||
Search is naive substring match; ranking is keyword-overlap. Real
|
||||
runtime implementations replace this with embeddings/dep-graph search.
|
||||
Policy enforcement (:class:`WorkspaceAccess`) IS applied here so tests
|
||||
cover the rejection paths.
|
||||
"""
|
||||
|
||||
_EXT_TO_TYPE: dict[str, FileType] = {
|
||||
".py": FileType.PYTHON,
|
||||
".ts": FileType.TYPESCRIPT,
|
||||
".tsx": FileType.TYPESCRIPT,
|
||||
".js": FileType.JAVASCRIPT,
|
||||
".jsx": FileType.JAVASCRIPT,
|
||||
".yaml": FileType.YAML,
|
||||
".yml": FileType.YAML,
|
||||
".json": FileType.JSON,
|
||||
".toml": FileType.TOML,
|
||||
".md": FileType.MARKDOWN,
|
||||
".sql": FileType.SQL,
|
||||
".sh": FileType.SHELL,
|
||||
}
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
files: dict[str, bytes],
|
||||
*,
|
||||
access: WorkspaceAccess,
|
||||
) -> None:
|
||||
self._files: dict[str, bytes] = dict(files)
|
||||
self._access = access
|
||||
self._grants: dict[str, WorkspaceGrant] = {}
|
||||
self._counter = 0
|
||||
|
||||
def _detect(self, path: str) -> FileType:
|
||||
for ext, ft in self._EXT_TO_TYPE.items():
|
||||
if path.endswith(ext):
|
||||
return ft
|
||||
return FileType.OTHER
|
||||
|
||||
def _denied(self, path: str) -> bool:
|
||||
for pat in self._access.deny_patterns:
|
||||
regex = re.compile(_glob_to_regex(pat))
|
||||
if regex.fullmatch(path):
|
||||
return True
|
||||
return False
|
||||
|
||||
async def search(
|
||||
self,
|
||||
*,
|
||||
query: str,
|
||||
types: Sequence[FileType] = (),
|
||||
limit: int = 20,
|
||||
) -> list[FileMatch]:
|
||||
if not self._access.enabled:
|
||||
raise WorkspaceDenied("workspace disabled by policy")
|
||||
type_set = set(types)
|
||||
terms = [t.lower() for t in re.split(r"\W+", query) if t]
|
||||
out: list[FileMatch] = []
|
||||
for path, data in self._files.items():
|
||||
if self._denied(path):
|
||||
continue
|
||||
ft = self._detect(path)
|
||||
if type_set and ft not in type_set:
|
||||
continue
|
||||
haystack = (path + "\n" + data.decode("utf-8", errors="ignore")).lower()
|
||||
score = sum(haystack.count(t) for t in terms)
|
||||
if score == 0:
|
||||
continue
|
||||
out.append(
|
||||
FileMatch(
|
||||
path=path,
|
||||
file_type=ft,
|
||||
score=float(score),
|
||||
size_bytes=len(data),
|
||||
)
|
||||
)
|
||||
out.sort(key=lambda m: -m.score)
|
||||
return out[:limit]
|
||||
|
||||
async def request_access(
|
||||
self,
|
||||
*,
|
||||
files: Sequence[FileMatch | str],
|
||||
mode: WorkspaceMode,
|
||||
reason: str,
|
||||
purpose: str = "",
|
||||
) -> WorkspaceGrant:
|
||||
if not self._access.enabled:
|
||||
raise WorkspaceDenied("workspace disabled by policy")
|
||||
if mode not in self._access.allowed_modes:
|
||||
raise WorkspaceDenied(f"mode {mode.value!r} not in allowed_modes")
|
||||
if self._access.require_reason and not reason.strip():
|
||||
raise WorkspaceDenied("reason required by policy")
|
||||
normalized: list[FileMatch] = []
|
||||
for f in files:
|
||||
if isinstance(f, str):
|
||||
if f not in self._files:
|
||||
raise WorkspaceDenied(f"unknown path: {f!r}")
|
||||
normalized.append(
|
||||
FileMatch(
|
||||
path=f,
|
||||
file_type=self._detect(f),
|
||||
size_bytes=len(self._files[f]),
|
||||
)
|
||||
)
|
||||
else:
|
||||
if f.path not in self._files:
|
||||
raise WorkspaceDenied(f"unknown path: {f.path!r}")
|
||||
normalized.append(f)
|
||||
if len(normalized) > self._access.max_files:
|
||||
raise WorkspaceDenied(
|
||||
f"requested {len(normalized)} files, max_files={self._access.max_files}"
|
||||
)
|
||||
for m in normalized:
|
||||
if self._denied(m.path):
|
||||
raise WorkspaceDenied(f"path denied by policy: {m.path}")
|
||||
total = sum(m.size_bytes for m in normalized)
|
||||
if total > self._access.max_total_size_bytes:
|
||||
raise WorkspaceDenied(
|
||||
f"total size {total} exceeds max_total_size_bytes"
|
||||
)
|
||||
self._counter += 1
|
||||
grant = WorkspaceGrant(
|
||||
grant_id=f"grant-{self._counter}",
|
||||
purpose=purpose,
|
||||
files=tuple(normalized),
|
||||
mode=mode,
|
||||
reason=reason,
|
||||
requires_human_approval=self._access.require_human_approval,
|
||||
)
|
||||
self._grants[grant.grant_id] = grant
|
||||
return grant
|
||||
|
||||
async def open_view(
|
||||
self,
|
||||
*,
|
||||
purpose: str,
|
||||
hints: Sequence[str] = (),
|
||||
file_types: Sequence[FileType] = (),
|
||||
max_files: int = 10,
|
||||
mode: WorkspaceMode = WorkspaceMode.READ_ONLY,
|
||||
reason: str | None = None,
|
||||
) -> WorkspaceView:
|
||||
query = " ".join([purpose, *hints])
|
||||
matches = await self.search(query=query, types=file_types, limit=max_files * 3)
|
||||
chosen = matches[:max_files]
|
||||
grant = await self.request_access(
|
||||
files=chosen,
|
||||
mode=mode,
|
||||
reason=reason or purpose,
|
||||
purpose=purpose,
|
||||
)
|
||||
return LocalWorkspaceView(grant, self)
|
||||
|
||||
async def list_grants(self) -> list[WorkspaceGrant]:
|
||||
return list(self._grants.values())
|
||||
|
||||
|
||||
def _glob_to_regex(pattern: str) -> str:
|
||||
"""Translate a simple ``foo/**/*.py`` style glob to a regex."""
|
||||
out: list[str] = []
|
||||
i = 0
|
||||
while i < len(pattern):
|
||||
c = pattern[i]
|
||||
if c == "*":
|
||||
if i + 1 < len(pattern) and pattern[i + 1] == "*":
|
||||
out.append(".*")
|
||||
i += 2
|
||||
if i < len(pattern) and pattern[i] == "/":
|
||||
i += 1
|
||||
else:
|
||||
out.append("[^/]*")
|
||||
i += 1
|
||||
elif c == "?":
|
||||
out.append("[^/]")
|
||||
i += 1
|
||||
elif c == ".":
|
||||
out.append(r"\.")
|
||||
i += 1
|
||||
else:
|
||||
out.append(re.escape(c))
|
||||
i += 1
|
||||
return "".join(out)
|
||||
Reference in New Issue
Block a user