strukto-ai--mirage/python/mirage/workspace/workspace.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 asyncio
import builtins
import logging
import sys
from collections.abc import AsyncIterator
from functools import partial
from typing import Any

from mirage.cache.file.config import CacheConfig, RedisCacheConfig
from mirage.cache.file.ram import RAMFileCacheStore
from mirage.cache.index import IndexConfig
from mirage.commands.builtin.utils.safeguard import (CommandTimeoutError,
                                                     run_with_timeout)
from mirage.commands.errors import FindParseError, UsageError
from mirage.commands.safeguard import resolve_safeguard

try:
    from mirage.cache.file.redis import RedisFileCacheStore
except ImportError:
    RedisFileCacheStore = None  # type: ignore[misc, assignment]
from mirage.io import IOResult
from mirage.observe.context import RecordingScope
from mirage.observe.observer import Observer
from mirage.observe.record import OpRecord
from mirage.observe.store import ObserverStore
from mirage.ops import Ops
from mirage.ops.open import make_open
from mirage.ops.os_patch import make_os_module
from mirage.provision import ProvisionResult
from mirage.resource.base import BaseResource
from mirage.resource.history import HISTORY_PREFIX, HistoryViewResource
from mirage.resource.ram import RAMResource
from mirage.shell.job_table import JobTable
from mirage.shell.parse import find_syntax_error, parse
from mirage.types import (DEFAULT_AGENT_ID, DEFAULT_SESSION_ID,
                          ConsistencyPolicy, DriftPolicy, FileStat, MountMode,
                          PathSpec, StateKey)
from mirage.utils.errors import format_fs_error
from mirage.workspace.abort import MirageAbortError
from mirage.workspace.dispatcher import Dispatcher
from mirage.workspace.file_prompt import build_file_prompt
from mirage.workspace.fuse import FuseManager
from mirage.workspace.mount import MountEntry, MountRegistry
from mirage.workspace.mount.namespace import Namespace
from mirage.workspace.mount.spec import Mount
from mirage.workspace.node import provision_node, run_command_tree
from mirage.workspace.session import (Session, SessionManager,
                                      reset_current_session,
                                      set_current_session)
from mirage.workspace.snapshot import (ContentDriftError, apply_state_dict,
                                       build_mount_args, check_drift,
                                       install_fingerprints, norm_mount_prefix,
                                       read_tar, requires_resource_override)
from mirage.workspace.snapshot import snapshot as _write_snapshot
from mirage.workspace.snapshot import to_state_dict

logger = logging.getLogger(__name__)


class Workspace:
    """Unified virtual filesystem over heterogeneous resources.

    Manages mounts, caching, and command execution.
    All ops are forwarded directly to the resolved resource.
    """

    def __init__(
        self,
        resources: dict[str, BaseResource | tuple | Mount],
        cache_limit: str | int = "512MB",
        cache: CacheConfig | None = None,
        index: IndexConfig | None = None,
        mode: MountMode = MountMode.READ,
        consistency: ConsistencyPolicy = ConsistencyPolicy.LAZY,
        session_id: str = DEFAULT_SESSION_ID,
        agent_id: str = DEFAULT_AGENT_ID,
        observe: ObserverStore | None = None,
    ) -> None:
        self._registry = MountRegistry()
        if isinstance(cache, RedisCacheConfig):
            if RedisFileCacheStore is None:
                raise ImportError(
                    "RedisCacheConfig requires the 'redis' extra. "
                    "Install with: pip install mirage-ai[redis]")
            self._cache = RedisFileCacheStore(
                cache_limit=cache.limit,
                url=cache.url,
                key_prefix=cache.key_prefix,
                max_drain_bytes=cache.max_drain_bytes,
            )
        else:
            limit = cache.limit if cache is not None else cache_limit
            max_drain = cache.max_drain_bytes if cache is not None else None
            self._cache = RAMFileCacheStore(cache_limit=limit,
                                            max_drain_bytes=max_drain)
        self._locked_paths: set[str] = set()
        self._closed = False
        self._drift_policy: DriftPolicy = DriftPolicy.OFF
        self._drift_check_pending: bool = False
        # Queued at Workspace.load: (mount, path, expected_fingerprint).
        # First dispatch/execute drains via asyncio.gather, then clears.
        self._pending_drift: list[tuple[MountEntry, str, str]] = []
        self.job_table = JobTable()
        self._current_agent_id: str = agent_id
        self._default_session_id = session_id
        self._default_agent_id = agent_id
        self._session_mgr = SessionManager(session_id)
        self._consistency = consistency
        self._registry.set_consistency(consistency)
        self._registry.attach_file_cache(self._cache)
        self._namespace = Namespace(self._registry)
        self._dispatcher = Dispatcher(self._namespace, self._cache,
                                      consistency)

        fuse_targets: list[tuple[str, bool | str]] = []
        for prefix, value in resources.items():
            mount_safeguards: dict = {}
            mount_fuse: bool | str = False
            if isinstance(value, Mount):
                prov = value.resource
                mount_mode = value.mode if value.mode is not None else mode
                if value.command_safeguards:
                    mount_safeguards = dict(value.command_safeguards)
                mount_fuse = value.fuse
            elif isinstance(value, tuple) and len(value) >= 2:
                prov = value[0]
                mount_mode = value[1]
                if len(value) >= 3 and value[2]:
                    mount_safeguards = dict(value[2])
            else:
                prov = value
                mount_mode = mode
            if index is not None:
                prov.set_index(index)
            mount_obj = self._registry.mount(prefix, prov, mount_mode)
            if mount_safeguards:
                mount_obj.command_safeguards.update(mount_safeguards)
            if mount_fuse:
                fuse_targets.append((prefix, mount_fuse))

        if self._registry.root_mount is None:
            self._registry.mount("/", RAMResource(), mode)

        self._fuse_mountpoints: dict[str, str] = {}
        self._fuse_managers: dict[str, FuseManager] = {}

        self.observer = Observer(store=observe)
        self._registry.mount(HISTORY_PREFIX,
                             HistoryViewResource(self.observer),
                             MountMode.READ)

        self._ops = Ops(self._registry.ops_mounts(),
                        on_write=self._invalidate_after_write_by_path,
                        observer=self.observer,
                        agent_id=agent_id,
                        session_id=session_id)

        for prefix, fuse_target in fuse_targets:
            mountpoint = fuse_target if isinstance(fuse_target, str) else None
            self.add_fuse_mount(prefix, mountpoint)

    async def history(self) -> list[dict]:
        """Command events recorded by the hidden recorder.

        Returns:
            list[dict]: All sessions' command events, timestamp order.
        """
        return await self.observer.command_events()

    @property
    def ops(self) -> Ops:
        return self._ops

    @property
    def namespace(self) -> Namespace:
        return self._namespace

    @property
    def cache(self):
        return self._cache

    @property
    def max_drain_bytes(self) -> int | None:
        return self._cache.max_drain_bytes

    @max_drain_bytes.setter
    def max_drain_bytes(self, value: int | None) -> None:
        self._cache.max_drain_bytes = value

    def mounts(self) -> list:
        return self._registry.mounts()

    @property
    def revisions(self) -> dict[str, str]:
        """Flat view of every mount's installed revision pins.

        Derived (read-only) — the source of truth lives per-mount on
        ``mount.revisions``. Useful for tests, audit ("which paths got
        pinned at load?"), and debugging. Empty until a snapshot is
        loaded with revisions in its manifest.
        """
        out: dict[str, str] = {}
        for m in self._registry.mounts():
            if m.revisions:
                out.update(m.revisions)
        return out

    def mount(self, prefix: str):
        return self._registry.mount_for(prefix)

    async def unmount(self, prefix: str) -> None:
        if self._closed:
            raise RuntimeError("Workspace is closed")
        stripped = prefix.strip("/")
        norm = ("/" + stripped + "/" if stripped else "/")
        if norm == "/":
            raise ValueError(f"cannot unmount the virtual root: {prefix!r}")
        if norm == "/dev/":
            raise ValueError("cannot unmount reserved prefix: '/dev/'")
        if norm == HISTORY_PREFIX + "/":
            raise ValueError(f"cannot unmount history view: "
                             f"{HISTORY_PREFIX!r}")
        removed = self._registry.unmount(prefix)
        self._ops.unmount(prefix)
        remaining = self._registry.mounts()
        still_instance = any(m.resource is removed.resource for m in remaining)
        still_kind = any(m.resource.name == removed.resource.name
                         for m in remaining)
        if not still_kind:
            self._ops._registry.unregister_resource(removed.resource.name)
        if not still_instance:
            close = getattr(removed.resource, "close", None)
            if callable(close):
                result = close()
                if hasattr(result, "__await__"):
                    await result

    def _register_fuse(self, prefix: str, mountpoint: str) -> None:
        for other_prefix, other_mp in self._fuse_mountpoints.items():
            if other_mp == mountpoint and other_prefix != prefix:
                raise ValueError(
                    f"FUSE mountpoint {mountpoint!r} already used by "
                    f"prefix {other_prefix!r}; mounts need distinct paths")
        self._fuse_mountpoints[prefix] = mountpoint

    def _deregister_fuse(self, prefix: str) -> None:
        self._fuse_mountpoints.pop(prefix, None)

    def add_fuse_mount(self,
                       prefix: str,
                       mountpoint: str | None = None) -> str:
        # Register a pinned path BEFORE mounting so a collision is rejected
        # without leaving a partial mount. Each mount gets its own manager,
        # so a workspace can expose any number of FUSE subtrees at once.
        if mountpoint is not None:
            self._register_fuse(prefix, mountpoint)
        fm = FuseManager()
        self._fuse_managers[prefix] = fm
        try:
            mp = fm.setup(self._ops, prefix, mountpoint)
        except Exception:
            # The mount never came up; drop the manager and any registered
            # path so fuse_mountpoints does not misreport it as live.
            self._fuse_managers.pop(prefix, None)
            self._deregister_fuse(prefix)
            raise
        if mountpoint is None:
            self._register_fuse(prefix, mp)
        return mp

    def remove_fuse_mount(self, prefix: str) -> None:
        fm = self._fuse_managers.pop(prefix, None)
        if fm is not None:
            fm.unmount()
        self._deregister_fuse(prefix)

    @property
    def fuse_mountpoint(self) -> str | None:
        if not self._fuse_mountpoints:
            return None
        if len(self._fuse_mountpoints) > 1:
            raise RuntimeError(
                "multiple FUSE mounts active; use fuse_mountpoints to "
                "select one by prefix")
        return next(iter(self._fuse_mountpoints.values()))

    @property
    def fuse_mountpoints(self) -> dict[str, str]:
        return dict(self._fuse_mountpoints)

    @property
    def _cwd(self) -> str:
        return self._session_mgr.cwd

    @_cwd.setter
    def _cwd(self, value: str) -> None:
        self._session_mgr.cwd = value

    @property
    def env(self) -> dict[str, str]:
        return self._session_mgr.env

    @env.setter
    def env(self, value: dict[str, str]) -> None:
        self._session_mgr.env = value

    @property
    def file_prompt(self) -> str:
        return build_file_prompt(self._registry.mounts())

    # ── lifecycle ───────────────────────────────────────────────────────────

    def __enter__(self) -> "Workspace":
        self._original_open = builtins.open
        self._original_os = sys.modules["os"]
        builtins.open = make_open(self._ops)
        sys.modules["os"] = make_os_module(self._ops)
        return self

    def __exit__(self, *_: object) -> None:
        builtins.open = self._original_open
        sys.modules["os"] = self._original_os
        self._close_parts()

    def _close_parts(self) -> None:
        if self._closed:
            return
        self._closed = True
        for fm in list(self._fuse_managers.values()):
            fm.unmount()
        self._fuse_managers.clear()
        self._fuse_mountpoints.clear()
        for job in self.job_table.running_jobs():
            self.job_table.kill(job.id)
        for task in self._cache._drain_tasks.values():
            task.cancel()
        self._cache._drain_tasks.clear()

    async def close(self) -> None:
        drain_tasks = list(self._cache._drain_tasks.values())
        self._close_parts()
        for task in drain_tasks:
            try:
                await task
            except asyncio.CancelledError:
                pass
        await self._cache.clear()

    # ── snapshot / load / copy ─────────────────────────────────────────────

    async def snapshot(self, target, *, compress: str | None = None) -> None:
        """Serialize this workspace to a tar.

        Captured:
            * Mount configs, sessions, history, finished jobs.
            * Cache bytes for fast replay.
            * One fingerprint entry per remote read (ETag-equivalent,
              plus a backend-specific ``revision`` when the resource
              exposes one — e.g. S3 ``VersionId``).

        NOT captured:
            * Live state of mounts with ``SUPPORTS_SNAPSHOT=False``
              (Gmail, Slack, Linear, etc.). Load logs a warning naming
              them.
            * Files the agent never touched.
            * Bytes of remote objects. Recovery of original bytes works
              only when the resource accepts a revision pin (S3 family
              today) and the recorded revision still exists on the
              source.

        Async because fingerprint capture stats each touched path on a
        ``SUPPORTS_SNAPSHOT`` mount.

        Args:
            target: filesystem path OR a writable file-like object.
            compress: None | "gz" | "bz2" | "xz".
        """
        await _write_snapshot(self, target, compress=compress)

    @classmethod
    async def load(
            cls,
            source,
            *,
            resources: dict | None = None,
            drift_policy: DriftPolicy = DriftPolicy.STRICT) -> "Workspace":
        """Reconstruct a Workspace from a tar.

        For every recorded read:

        1. If the manifest entry carries a ``revision`` (e.g. S3
           ``VersionId``), the load installs it into the owning
           ``mount.revisions``. Replay reads pin to that revision via
           the ``revision_for`` contextvar lookup, so the original
           bytes are served. Drift check is skipped for these paths —
           the pin guarantees bytes match by construction.
        2. If the entry carries only a ``fingerprint`` (no stable
           revision), the load queues a drift check. STRICT raises
           ``ContentDriftError`` on the first mismatch; OFF skips the
           check entirely and evicts the snapshot cache so reads serve
           current state.

        Drift check is eager (fires once on the first dispatch or
        execute), so downstream code can rely on consistent state.

        Args:
            source: filesystem path OR a readable file-like object.
            resources: {prefix: Resource} overrides for mounts saved
                with redacted creds.
            drift_policy: STRICT (default) raises on mismatch. OFF
                disables drift checking and evicts snapshot cache for
                fingerprinted paths.
        """
        return await cls.from_state(read_tar(source),
                                    resources=resources,
                                    drift_policy=drift_policy)

    @classmethod
    async def from_state(
            cls,
            state: dict,
            *,
            resources: dict | None = None,
            drift_policy: DriftPolicy = DriftPolicy.STRICT) -> "Workspace":
        """Reconstruct a Workspace directly from a state dict (no tar).

        The in-process inverse of ``to_state_dict``: build the mounts,
        restore content/cache/history, then install drift fingerprints.
        ``load`` is this plus a tar read; callers that already hold a
        state dict (e.g. a version checkout) should use this and skip the
        tar round-trip.

        Args:
            state: a state dict from ``to_state_dict`` or a version.
            resources: {prefix: Resource} overrides for mounts saved
                with redacted creds.
            drift_policy: STRICT (default) raises on mismatch. OFF
                disables drift checking and evicts snapshot cache for
                fingerprinted paths.
        """
        ws = await cls._from_state(state, resources=resources)
        install_fingerprints(ws,
                             state.get(StateKey.FINGERPRINTS) or [],
                             drift_policy)
        live_only = state.get(StateKey.LIVE_ONLY_MOUNTS) or []
        if live_only:
            logger.warning(
                "Workspace.from_state: %s mount(s) opt out of snapshot "
                "replay; reads against them will serve current state with "
                "no drift detection: %s", len(live_only), live_only)
        return ws

    async def copy(self) -> "Workspace":
        # Reuse this process's resources so remote backends (S3, Redis,
        # GDrive) stay shared between original and copy. Local backends
        # (RAM, Disk) restore their content fresh into the new resources
        # — see snapshot.api.snapshot docstring for the rationale.
        # Only reuse resources whose state has redacted secrets or connection
        # material. Local content resources (RAM, Disk) are reconstructed
        # fresh so the copy's writes don't clobber the original's data.
        state = await to_state_dict(self)
        auto_prefixes = {"/dev/", norm_mount_prefix(HISTORY_PREFIX)}
        prefix_to_resource = {
            m.prefix: m.resource
            for m in self._registry.mounts() if m.prefix not in auto_prefixes
        }
        resources = {
            m["prefix"]: prefix_to_resource[m["prefix"]]
            for m in state["mounts"] if requires_resource_override(m)
            and m["prefix"] in prefix_to_resource
        }
        return await type(self)._from_state(state, resources=resources)

    @classmethod
    async def _from_state(cls,
                          state: dict,
                          *,
                          resources: dict | None = None) -> "Workspace":
        args = build_mount_args(state, resources)
        ws = cls(args.mount_args,
                 consistency=args.consistency,
                 session_id=args.default_session_id,
                 agent_id=args.default_agent_id)
        await apply_state_dict(ws, state)
        return ws

    def __deepcopy__(self, memo) -> "Workspace":
        raise NotImplementedError(
            "Workspace.copy is async (it captures fingerprints for replay). "
            "Call `await ws.copy()` directly instead of `copy.deepcopy(ws)`.")

    def __copy__(self) -> "Workspace":
        raise NotImplementedError("Workspace has no useful shallow copy — "
                                  "use `await ws.copy()`.")

    # ── session lifecycle ──────────────────────────────────────────────────

    def create_session(
            self,
            session_id: str,
            allowed_mounts: frozenset[str] | None = None) -> Session:
        if allowed_mounts is not None:
            normalized = {("/" + m.strip("/")) for m in allowed_mounts}
            normalized.update(self._infrastructure_mount_prefixes())
            allowed_mounts = frozenset(normalized)
        return self._session_mgr.create(session_id,
                                        allowed_mounts=allowed_mounts)

    def _infrastructure_mount_prefixes(self) -> set[str]:
        """Mount prefixes a session is always allowed to touch.

        The virtual root (where text-processing commands like ``wc``
        without a path argument resolve), the device mount, and the
        history view are infrastructure: they hold no user
        credentials, and rejecting them would break common shell
        idioms or the history builtin.
        """
        prefixes = {"/dev", HISTORY_PREFIX}
        root_mount = self._registry.root_mount
        if root_mount is not None:
            prefixes.add("/" + root_mount.prefix.strip("/"))
        return prefixes

    def get_session(self, session_id: str) -> Session:
        return self._session_mgr.get(session_id)

    def list_sessions(self) -> list[Session]:
        return self._session_mgr.list()

    async def close_session(self, session_id: str) -> None:
        await self._session_mgr.close(session_id)

    async def close_all_sessions(self) -> None:
        await self._session_mgr.close_all()

    # ── mount management ────────────────────────────────────────────────────

    async def dispatch(self, op: str, path: PathSpec,
                       **kwargs: Any) -> tuple[Any, IOResult]:
        if self._drift_check_pending:
            await self._run_pending_drift_check()
        return await self._dispatcher.dispatch(op, path, **kwargs)

    async def _run_pending_drift_check(self) -> None:
        """Drain the post-load drift check.

        Called once on the first async entry point (``dispatch`` or
        ``execute``) after ``Workspace.load`` with a non-OFF drift
        policy. Stats every queued ``(mount, path, expected_fingerprint)``
        triple against the live source in parallel and raises
        :class:`ContentDriftError` on the first mismatch. Subsequent
        calls are no-ops.

        Pinned paths (those whose manifest entry carried a stable
        revision) are never enqueued, because the pin guarantees bytes
        match by construction.

        Stats are issued with ``asyncio.gather`` so first-op latency
        does not scale linearly with the number of recorded reads.
        """
        self._drift_check_pending = False
        if not self._pending_drift:
            return
        checks = [
            check_drift(self, path, fingerprint)
            for _, path, fingerprint in self._pending_drift
        ]
        self._pending_drift.clear()
        results = await asyncio.gather(*checks, return_exceptions=True)
        for r in results:
            if isinstance(r, BaseException):
                raise r

    async def stat(self, path: str) -> FileStat:
        scope = PathSpec(virtual=path,
                         directory=path,
                         resource_path="",
                         resolved=True)
        result, _ = await self.dispatch("stat", scope)
        return result

    async def readdir(self, path: str) -> list[str]:
        scope = PathSpec(virtual=path,
                         directory=path,
                         resource_path="",
                         resolved=False)
        raw, _ = await self.dispatch("readdir", scope)
        return raw

    # ── execution ────────────────────────────────────────────────────────────

    async def apply_io(self,
                       io: IOResult,
                       records: list[OpRecord] | None = None) -> None:
        await self._dispatcher.apply_io(io, records=records)

    async def _invalidate_after_write_by_path(self, path: str) -> None:
        await self._dispatcher.invalidate_after_write_by_path(path)

    def _session_cwd(self, session_id: str) -> str | None:
        try:
            return self._session_mgr.get(session_id).cwd
        except KeyError:
            return None

    async def _plan_eval_stub(self, cmd: str, **opts: Any) -> IOResult:
        """Inert evaluator for provision walks.

        A dry run must never execute: a command substitution with side
        effects ($(tee ...)) would otherwise run while "estimating".
        Substitutions expand to empty, so affected words degrade the
        plan to honest UNKNOWN instead of resolving via execution.
        """
        return IOResult()

    async def _exec_recursion(self, cancel: asyncio.Event | None, cmd: str,
                              **opts: Any) -> Any:
        # The executor's internal eval ($(), source, eval, xargs, ...):
        # never a typed line, so it must not record a history entry or
        # open its own recording context (GNU: history is appended by
        # the line reader, the evaluator can't touch it).
        return await self.execute(cmd, cancel=cancel, record=False, **opts)

    async def execute(
        self,
        command: str,
        session_id: str | None = None,
        stdin: AsyncIterator[bytes] | bytes | None = None,
        provision: bool = False,
        agent_id: str = DEFAULT_AGENT_ID,
        cwd: str | None = None,
        env: dict[str, str] | None = None,
        cancel: asyncio.Event | None = None,
        record: bool = True,
    ) -> IOResult | ProvisionResult:
        """Execute a shell command in the workspace.

        Args:
            command: The shell command string to execute.
            session_id: Session whose persistent state hosts the command.
            stdin: Optional stdin payload (bytes or async byte iterator).
            provision: If True, return a ProvisionResult instead of running.
            agent_id: Agent identifier for observability and history.
            cwd: Per-call working directory override. When provided, the
                command runs in an ephemeral session clone (bash subshell
                semantics): the persistent session's cwd is unchanged and
                any `cd` inside the command does not leak.
            env: Per-call environment overrides layered on top of the
                session's env. Like cwd, these apply only to an ephemeral
                clone, so `export` inside the command does not leak back
                to the persistent session.
            cancel: Optional asyncio.Event used to abort execution
                mid-flight. When set, the executor raises MirageAbortError
                at the next gate (entry to each node) and races inside
                blocking sleeps so cancellation is observed promptly.
            record: When False, run without logging a history entry or
                opening a recording context; ops emitted by the command
                flow into the caller's recorder. Used by the executor's
                internal evaluations and available to SDK callers that
                need an unrecorded run.
        """
        if cancel is not None and cancel.is_set():
            raise MirageAbortError()
        if self._drift_check_pending:
            await self._run_pending_drift_check()

        if session_id is None:
            session_id = self._session_mgr.default_id
        session = self._session_mgr.get(session_id)
        use_override = cwd is not None or env is not None
        if use_override:
            overrides: dict[str, Any] = {}
            if cwd is not None:
                overrides["cwd"] = cwd
            if env is not None:
                overrides["env"] = {**session.env, **env}
            effective_session = session.fork(**overrides)
        else:
            effective_session = session
        self._current_agent_id = agent_id
        io = IOResult()
        # The line-reader decision (GNU: history is appended where the
        # typed line is read, never inside the evaluator). Internal
        # evaluations and provision runs get an inert scope.
        is_line = record and not provision
        scope = RecordingScope(active=is_line)

        exec_recursion = partial(self._exec_recursion, cancel)

        session_token = set_current_session(effective_session)
        try:
            ast = parse(command)
            offending = find_syntax_error(ast)
            if offending is not None:
                snippet = offending.strip()[:40]
                err = (f"mirage: syntax error near {snippet!r}\n".encode()
                       if snippet else b"mirage: syntax error in command\n")
                io = IOResult(exit_code=2, stderr=err)
                return io
            if provision:
                prov_name = command.strip().split()[0] if command.strip(
                ) else None
                prov_resolved = (resolve_safeguard(prov_name)
                                 if prov_name else None)
                prov_timeout = (prov_resolved.timeout_seconds
                                if prov_resolved is not None else None)
                return await run_with_timeout(
                    provision_node(self._registry, self.dispatch,
                                   self._plan_eval_stub, self._namespace, ast,
                                   effective_session), prov_timeout, prov_name)
            io, _ = await run_command_tree(
                self.dispatch,
                self._registry,
                self._namespace,
                self.job_table,
                exec_recursion,
                self._current_agent_id,
                ast,
                effective_session,
                stdin,
                cancel,
            )
            session.last_exit_code = io.exit_code
            await self.apply_io(io, records=scope.records)
            return io
        except CommandTimeoutError as exc:
            logger.debug("command %r timed out after %ss", exc.command,
                         exc.seconds)
            if cancel is not None:
                cancel.set()
            msg = (str(exc) + "\n").encode()
            io = IOResult(exit_code=124, stderr=msg)
            session.last_exit_code = 124
            return io
        except (MirageAbortError, ContentDriftError):
            raise
        except FindParseError as exc:
            msg = f"{exc}\n".encode()
            io = IOResult(exit_code=1, stderr=msg)
            return io
        except UsageError as exc:
            msg = f"{exc}\n".encode()
            io = IOResult(exit_code=exc.exit_code, stderr=msg)
            return io
        except OSError as exc:
            cmd_name = command.split()[0] if command.split() else command
            msg = format_fs_error(cmd_name, exc)
            io = IOResult(exit_code=1, stderr=msg)
            return io
        except Exception as exc:
            io = IOResult(exit_code=1, stderr=str(exc).encode())
            return io
        finally:
            # One rule on every path: an op that happened is always
            # accounted, in byte accounting (which feeds snapshot
            # fingerprints/drift) and as observer op events. The
            # command event's exit_code says whether the line that
            # emitted them succeeded.
            scope.close()
            reset_current_session(session_token)
            self._ops.records.extend(scope.records)
            if is_line:
                await self.observer.log_execution(
                    command, io, scope.records, agent_id, session_id,
                    self._session_cwd(session_id))