bazaar-compute-node 0.1.3__py3-none-any.whl

bazaar_compute_node/__init__.py

@@ -0,0 +1,3 @@
+from importlib.metadata import version
+
+__version__ = version("bazaar-compute-node")

bazaar_compute_node/app/__init__.py

@@ -0,0 +1 @@
+"""Application composition and local command transport."""

bazaar_compute_node/app/application.py

@@ -0,0 +1,398 @@
+from __future__ import annotations
+
+import asyncio
+import hmac
+import os
+import re
+import secrets
+import signal
+from collections.abc import Mapping, Sequence
+from pathlib import Path
+
+from ..core.channel import ChannelContext, IChannel
+from ..core.lifecycle import TimeoutBudget
+from ..core.models import RuntimeSession
+from ..core.observability import IAudit
+from ..core.orchestration import SessionOrchestrator
+from ..core.paths import resolve_data_dir, resolve_workspace_dir
+from ..core.runtime import IRuntime, RuntimeCommandContext, RuntimeSandboxMode
+from ..core.storage import IStorage, NodeIdentity
+from .attachments import AttachmentMaterializer
+from .command import (
+    CommandDispatcher,
+    CommandDispatchError,
+)
+from .registry import AdapterFactories
+from .transport import LocalCommandServer
+from .wrapper import install_bcc_wrapper, remove_bcc_wrapper
+
+CommandRecord = tuple[str, tuple[str, ...]]
+
+_ENVIRONMENT_NAME = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$")
+_PLATFORM_ENVIRONMENT = {
+    "HOME",
+    "LANG",
+    "LC_ALL",
+    "PATH",
+    "TMPDIR",
+    "TEMP",
+    "TMP",
+    "SystemRoot",
+    "ComSpec",
+    "PATHEXT",
+    "USERPROFILE",
+}
+_FORBIDDEN_ENVIRONMENT = {
+    "BCN_WECOM_BOT_SECRET",
+    "DATABASE_URL",
+    "AWS_ACCESS_KEY_ID",
+    "AWS_SECRET_ACCESS_KEY",
+    "OPENAI_API_KEY",
+}
+
+
+class NodeApplication:
+    """Generic application lifecycle for one dynamically composed node."""
+
+    def __init__(
+        self,
+        *,
+        factories: AdapterFactories,
+        endpoint_path: Path | None = None,
+        node_id: str = "bcn-node",
+        workspace_id: str | None = None,
+        runtime_options: Mapping[str, str] | None = None,
+        runtime_sandbox_mode: RuntimeSandboxMode = RuntimeSandboxMode.WORKSPACE_WRITE,
+        runtime_network_access: bool = True,
+        channel_options: Mapping[str, object] | None = None,
+        runtime_environment_include: Sequence[str] = (),
+        timeout_budget: TimeoutBudget | None = None,
+    ) -> None:
+        self.data_dir = resolve_data_dir()
+        self.runtime_options = dict(runtime_options or {})
+        self.timeout_budget = timeout_budget or TimeoutBudget(
+            startup_seconds=60,
+            provider_call_seconds=600,
+            command_seconds=10,
+            shutdown_seconds=5,
+        )
+        self.storage: IStorage = factories.storage()
+        self.audit: IAudit = factories.audit()
+        self.command_log: list[CommandRecord] = []
+        self._wrapper_path: Path | None = None
+        self._identity: NodeIdentity | None = None
+        self._session_capabilities: dict[str, str] = {}
+        self._runtime_session_ids: dict[str, str] = {}
+        self._started = False
+        self._stopped = asyncio.Event()
+        self._runtime_environment_include = tuple(runtime_environment_include)
+        self._attachment_materializer = AttachmentMaterializer(
+            self._workspace_path, self._referenced_attachment_paths
+        )
+        self.channel: IChannel = factories.channel(
+            ChannelContext(
+                attachments=self._attachment_materializer,
+                options=dict(channel_options or {}),
+                workspace=self._workspace_path,
+            )
+        )
+        self._runtime_context = RuntimeCommandContext(
+            run_command=self._run_runtime_command,
+            environment_for_session=self._runtime_environment,
+            node_id=node_id,
+            runtime_options=self.runtime_options,
+            sandbox_mode=runtime_sandbox_mode,
+            network_access=runtime_network_access,
+            startup_timeout_seconds=self.timeout_budget.startup_seconds,
+        )
+        self.runtime: IRuntime = factories.runtime(self._runtime_context)
+        self.orchestrator = SessionOrchestrator(
+            node_id=node_id,
+            workspace_id=workspace_id,
+            channel=self.channel,
+            runtime=self.runtime,
+            storage=self.storage,
+            audit=self.audit,
+            timeout_budget=self.timeout_budget,
+            on_node_initialized=self._ensure_workspace,
+        )
+        self.command_service = self.orchestrator.command_service
+        control_handler = None
+        if factories.control is not None:
+            control_handler = factories.control(self._adapter_context())
+        self.command_dispatcher = CommandDispatcher(
+            self.command_service,
+            timeout_budget=self.timeout_budget,
+            control_handler=self._handle_control,
+            session_binding_validator=self._validate_session_binding,
+        )
+        self.command_server = LocalCommandServer(
+            self.command_dispatcher,
+            endpoint_path=endpoint_path,
+        )
+        self._provider_control_handler = control_handler
+
+    @property
+    def endpoint(self) -> str:
+        return self.command_server.endpoint
+
+    async def start(self) -> None:
+        if self._started:
+            return
+        self.data_dir.mkdir(parents=True, exist_ok=True)
+        if os.name != "nt":
+            self.data_dir.chmod(0o700)
+        self._wrapper_path = install_bcc_wrapper(self.data_dir / "bin")
+        try:
+            await self.command_server.start()
+            await self.orchestrator.start(
+                timeout=self.timeout_budget.startup_seconds,
+            )
+        except BaseException:
+            await self.stop()
+            raise
+        self._started = True
+        self._stopped.clear()
+        self.command_dispatcher.start_accepting()
+
+    async def stop(self) -> None:
+        if not self._started:
+            try:
+                await self.command_dispatcher.drain(
+                    timeout=self.timeout_budget.shutdown_seconds,
+                )
+            finally:
+                try:
+                    await self.command_server.stop()
+                finally:
+                    self._cleanup_bcc_wrapper()
+            return
+        self._started = False
+        self.command_dispatcher.stop_accepting()
+        try:
+            await self.command_dispatcher.drain(
+                timeout=self.timeout_budget.shutdown_seconds,
+            )
+        finally:
+            try:
+                await self.orchestrator.stop(
+                    timeout=self.timeout_budget.shutdown_seconds,
+                )
+            finally:
+                try:
+                    await self.command_server.stop()
+                finally:
+                    try:
+                        self._cleanup_bcc_wrapper()
+                    finally:
+                        self._stopped.set()
+
+    def _cleanup_bcc_wrapper(self) -> None:
+        wrapper_path = self._wrapper_path
+        if wrapper_path is None:
+            return
+        remove_bcc_wrapper(wrapper_path)
+        self._wrapper_path = None
+
+    async def wait(self) -> None:
+        loop = asyncio.get_running_loop()
+        for signum in (signal.SIGINT, signal.SIGTERM):
+            try:
+                loop.add_signal_handler(signum, self._stopped.set)
+            except NotImplementedError, RuntimeError:
+                pass
+        await self._stopped.wait()
+
+    async def _ensure_workspace(self, identity: NodeIdentity) -> None:
+        self._identity = identity
+        workspace_dir = resolve_workspace_dir(identity.workspace_id)
+        await asyncio.to_thread(
+            workspace_dir.mkdir,
+            parents=True,
+            exist_ok=True,
+            mode=0o700,
+        )
+        if os.name != "nt":
+            await asyncio.to_thread(workspace_dir.chmod, 0o700)
+        await self._attachment_materializer.reconcile()
+
+    def _workspace_path(self) -> Path:
+        identity = self._identity
+        if identity is None:
+            raise RuntimeError("node identity has not been initialized")
+        return resolve_workspace_dir(identity.workspace_id)
+
+    async def _referenced_attachment_paths(self) -> set[str]:
+        async with self.storage.transaction() as transaction:
+            return set(await transaction.list_ready_attachment_paths())
+
+    def _adapter_context(self) -> Mapping[str, object]:
+        return {
+            "channel": self.channel,
+            "runtime": self.runtime,
+            "storage": self.storage,
+            "audit": self.audit,
+            "command_log": self.command_log,
+            "is_started": lambda: self._started,
+        }
+
+    async def _handle_control(
+        self, request: Mapping[str, object]
+    ) -> Mapping[str, object]:
+        if request.get("operation") == "health":
+            identity = self._identity
+            return {
+                "started": self._started,
+                "accepting": self.command_dispatcher.accepting,
+                "channel": self.channel.name,
+                "channel_health": dict(self.channel.health),
+                "runtime": self.runtime.name,
+                "storage": self.storage.name,
+                "audit": self.audit.name,
+                "node_id": identity.node_id if identity is not None else None,
+                "workspace_id": (
+                    identity.workspace_id if identity is not None else None
+                ),
+            }
+        if request.get("operation") == "shutdown":
+            self._stopped.set()
+            return {"accepted": True, "operation": "shutdown"}
+        if self._provider_control_handler is None:
+            raise ValueError("control operation is not supported")
+        return await self._provider_control_handler(request)
+
+    async def _validate_session_binding(
+        self,
+        session_id: str,
+        request: Mapping[str, object],
+    ) -> None:
+        runtime_session_id = request.get("runtime_session_id")
+        session_capability = request.get("session_capability")
+        async with self.storage.transaction() as transaction:
+            bcn_session = await transaction.get_bcn_session(session_id)
+            if bcn_session is None:
+                raise CommandDispatchError(
+                    "SESSION_NOT_FOUND",
+                    f"unknown bcn session: {session_id}",
+                )
+            runtime_session = await transaction.find_runtime_session(session_id)
+        if runtime_session is None or runtime_session_id != (runtime_session.id):
+            raise CommandDispatchError(
+                "SESSION_BINDING_FAILED",
+                "runtime session binding is invalid",
+            )
+        expected_capability = self._session_capabilities.get(session_id)
+        if (
+            expected_capability is None
+            or not isinstance(session_capability, str)
+            or not hmac.compare_digest(
+                session_capability.encode(), expected_capability.encode()
+            )
+        ):
+            raise CommandDispatchError(
+                "SESSION_BINDING_FAILED",
+                "session capability is invalid",
+            )
+
+    def _runtime_environment(self, session: RuntimeSession) -> Mapping[str, str]:
+        self._runtime_session_ids[session.bcn_session_id] = session.id
+        return self._build_command_environment(
+            session.bcn_session_id,
+            session.id,
+        )
+
+    def _build_command_environment(
+        self,
+        session_id: str,
+        runtime_session_id: str,
+    ) -> dict[str, str]:
+        if not session_id:
+            raise ValueError("session_id must be a non-empty string")
+        if not runtime_session_id:
+            raise ValueError("runtime_session_id must be a non-empty string")
+        wrapper_path = self._wrapper_path
+        if wrapper_path is None:
+            raise RuntimeError("bcc wrapper is not installed")
+        session_capability = self._session_capabilities.setdefault(
+            session_id,
+            secrets.token_urlsafe(32),
+        )
+        wrapper_directory = str(wrapper_path.parent)
+        allowed = set(_PLATFORM_ENVIRONMENT)
+        for name in self.runtime.environment_variable_names():
+            if not _ENVIRONMENT_NAME.fullmatch(name):
+                raise ValueError(f"runtime environment name is invalid: {name}")
+            if name.startswith("BCN_") or name in _FORBIDDEN_ENVIRONMENT:
+                raise ValueError(f"runtime environment name is reserved: {name}")
+            allowed.add(name)
+        for name in self._runtime_environment_include:
+            if not _ENVIRONMENT_NAME.fullmatch(name):
+                raise ValueError(f"runtime environment name is invalid: {name}")
+            if name.startswith("BCN_") or name in _FORBIDDEN_ENVIRONMENT:
+                raise ValueError(f"runtime environment name is reserved: {name}")
+            if name not in os.environ:
+                raise ValueError(f"runtime environment variable is missing: {name}")
+            allowed.add(name)
+        environment = {
+            name: os.environ[name] for name in sorted(allowed) if os.environ.get(name)
+        }
+        environment["PATH"] = os.pathsep.join(
+            (wrapper_directory, environment.get("PATH", os.defpath))
+        )
+        environment.update(
+            {
+                "BCN_ENDPOINT": self.endpoint,
+                "BCN_SESSION_ID": session_id,
+                "BCN_RUNTIME_SESSION_ID": runtime_session_id,
+                "BCN_COMMAND_CAPABILITY": session_capability,
+            }
+        )
+        return environment
+
+    async def _run_runtime_command(
+        self,
+        session_id: str,
+        arguments: Sequence[str],
+        body: str | None,
+    ) -> None:
+        if not session_id:
+            raise ValueError("session_id must be a non-empty string")
+        if not arguments or any(
+            not isinstance(argument, str) for argument in arguments
+        ):
+            raise ValueError("runtime command arguments must be non-empty text")
+        wrapper_path = self._wrapper_path
+        if wrapper_path is None:
+            raise RuntimeError("bcc wrapper is not installed")
+        self.command_log.append((session_id, tuple(arguments)))
+        runtime_session_id = self._runtime_session_ids.get(
+            session_id,
+            f"runtime-{session_id}",
+        )
+        environment = self._build_command_environment(
+            session_id,
+            runtime_session_id,
+        )
+        process = await asyncio.create_subprocess_exec(
+            str(wrapper_path),
+            *arguments,
+            stdin=asyncio.subprocess.PIPE,
+            stdout=asyncio.subprocess.PIPE,
+            stderr=asyncio.subprocess.PIPE,
+            env=environment,
+        )
+        input_data = body.encode() if body is not None else None
+        try:
+            _stdout, stderr = await process.communicate(input=input_data)
+        except asyncio.CancelledError:
+            if process.returncode is None:
+                process.terminate()
+                await process.wait()
+            raise
+        if process.returncode != 0:
+            error = stderr.decode(errors="replace").strip()
+            command = " ".join(arguments)
+            raise RuntimeError(f"bcc command failed ({command}): {error}")
+
+
+__all__ = ["CommandRecord", "NodeApplication"]

bazaar_compute_node/app/attachments.py

@@ -0,0 +1,154 @@
+from __future__ import annotations
+
+import asyncio
+import mimetypes
+import os
+import re
+import shutil
+from collections.abc import AsyncIterable, Awaitable, Callable
+from pathlib import Path, PurePosixPath
+from uuid import UUID, uuid7
+
+from ..core.models import InboundAttachment
+
+_SAFE_SUFFIX = re.compile(r"^\.[A-Za-z0-9]{1,10}$")
+
+
+class AttachmentMaterializer:
+    def __init__(
+        self,
+        workspace: Callable[[], Path],
+        referenced_paths: Callable[[], Awaitable[set[str]]],
+        *,
+        max_file_bytes: int = 25 * 1024 * 1024,
+        max_workspace_bytes: int = 1024 * 1024 * 1024,
+    ) -> None:
+        self._workspace = workspace
+        self._referenced_paths = referenced_paths
+        self._max_file_bytes = max_file_bytes
+        self._max_workspace_bytes = max_workspace_bytes
+        self._lock = asyncio.Lock()
+
+    async def reconcile(self) -> None:
+        root = self._workspace() / "attachments"
+        staging = root / ".staging"
+        await asyncio.to_thread(staging.mkdir, parents=True, exist_ok=True, mode=0o700)
+        for path in await asyncio.to_thread(lambda: tuple(staging.iterdir())):
+            if path.is_file() and not path.is_symlink():
+                await asyncio.to_thread(path.unlink)
+        referenced = await self._referenced_paths()
+        for path in await asyncio.to_thread(lambda: tuple(root.iterdir())):
+            if path == staging or not path.is_dir() or path.is_symlink():
+                continue
+            try:
+                if UUID(path.name).version != 7:
+                    continue
+            except ValueError:
+                continue
+            prefix = str(PurePosixPath("attachments", path.name)) + "/"
+            if any(relative_path.startswith(prefix) for relative_path in referenced):
+                continue
+            await asyncio.to_thread(shutil.rmtree, path)
+
+    async def materialize(
+        self,
+        source: bytes | AsyncIterable[bytes],
+        *,
+        name: str,
+        kind: str,
+        media_type: str | None = None,
+    ) -> InboundAttachment:
+        if not name or not kind:
+            raise ValueError("attachment name and kind must be non-empty")
+        attachment_id = str(uuid7())
+        suffix = Path(name).suffix
+        if not _SAFE_SUFFIX.fullmatch(suffix):
+            guessed = mimetypes.guess_extension(media_type or "") or ".bin"
+            suffix = guessed if _SAFE_SUFFIX.fullmatch(guessed) else ".bin"
+        relative = PurePosixPath(
+            "attachments", attachment_id, f"content{suffix.lower()}"
+        )
+        root = self._workspace() / "attachments"
+        staging = root / ".staging"
+        destination = self._workspace().joinpath(*relative.parts)
+        temporary = staging / f"{attachment_id}.part"
+        async with self._lock:
+            await asyncio.to_thread(
+                staging.mkdir, parents=True, exist_ok=True, mode=0o700
+            )
+            current_size = await asyncio.to_thread(self._stored_size, root)
+            size = 0
+            try:
+                with temporary.open("xb") as output:
+                    if isinstance(source, bytes):
+                        size = len(source)
+                        if size > self._max_file_bytes:
+                            raise ValueError(
+                                "attachment exceeds the per-file size limit"
+                            )
+                        if current_size + size > self._max_workspace_bytes:
+                            raise ValueError("attachment workspace quota exceeded")
+                        await asyncio.to_thread(output.write, source)
+                    else:
+                        async for chunk in source:
+                            if not isinstance(chunk, bytes):
+                                raise TypeError("attachment stream must yield bytes")
+                            size += len(chunk)
+                            if size > self._max_file_bytes:
+                                raise ValueError(
+                                    "attachment exceeds the per-file size limit"
+                                )
+                            if current_size + size > self._max_workspace_bytes:
+                                raise ValueError("attachment workspace quota exceeded")
+                            await asyncio.to_thread(output.write, chunk)
+                    await asyncio.to_thread(output.flush)
+                    await asyncio.to_thread(os.fsync, output.fileno())
+                await asyncio.to_thread(
+                    destination.parent.mkdir, parents=True, mode=0o700
+                )
+                await asyncio.to_thread(os.replace, temporary, destination)
+                if os.name != "nt":
+                    await asyncio.to_thread(destination.chmod, 0o600)
+            except BaseException:
+                if temporary.exists():
+                    await asyncio.to_thread(temporary.unlink)
+                raise
+        return InboundAttachment(
+            attachment_id=attachment_id,
+            name=Path(name).name,
+            kind=kind,
+            state="ready",
+            media_type=media_type,
+            relative_path=str(relative),
+            size_bytes=size,
+        )
+
+    def failed(
+        self,
+        *,
+        name: str,
+        kind: str,
+        error: str,
+        media_type: str | None = None,
+    ) -> InboundAttachment:
+        return InboundAttachment(
+            attachment_id=str(uuid7()),
+            name=Path(name).name or "attachment.bin",
+            kind=kind,
+            state="failed",
+            media_type=media_type,
+            error=error,
+        )
+
+    @staticmethod
+    def _stored_size(root: Path) -> int:
+        if not root.exists():
+            return 0
+        return sum(
+            path.stat().st_size
+            for path in root.glob("*/content.*")
+            if path.is_file() and not path.is_symlink()
+        )
+
+
+__all__ = ["AttachmentMaterializer"]