# ========= 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 dataclasses import inspect from typing import Any, Callable from mirage.cache.context import push_cache_manager from mirage.cache.manager import CacheManager from mirage.commands.builtin.utils.safeguard import (apply_op_safeguard, run_with_timeout) from mirage.commands.config import RegisteredCommand from mirage.commands.resolve import get_extension from mirage.commands.safeguard import CommandSafeguard, resolve_safeguard from mirage.commands.spec import CommandSpec from mirage.io.types import ByteSource, IOResult from mirage.observe.context import (push_mount_prefix, push_revisions, reset_revisions, with_mount_prefix, with_revisions) from mirage.ops.registry import RegisteredOp from mirage.resource.base import BaseResource from mirage.types import ConsistencyPolicy, MountMode, PathSpec from mirage.utils.key_prefix import mount_key def _wrap_cmd_streams( result: tuple[ByteSource | None, IOResult], mount_prefix: str, revisions: dict[str, str] | None, ) -> tuple[ByteSource | None, IOResult]: """Wrap any async-iterator streams in ``result`` with the mount prefix and active revisions, so ``record_stream`` and ``revision_for`` calls inside the lazy backend body see the right context when consumed after this frame exits. Mirrors the ``exit_on_empty`` pattern: thin async-gen wrapper that side-effects the recorder state as bytes flow through. Same object appearing in both the primary stream and IOResult.reads/writes is wrapped once (dedup by identity). Args: result: ``(stream, io)`` as returned by a command handler. mount_prefix: prefix to push during stream consumption. revisions: revisions map to push during stream consumption (None when the mount has no pins installed). """ stream, io = result seen: dict[int, ByteSource] = {} def _wrap(obj: ByteSource | None) -> ByteSource | None: if obj is None or isinstance(obj, (bytes, bytearray)): return obj oid = id(obj) if oid in seen: return seen[oid] wrapped = with_mount_prefix(mount_prefix, obj) if revisions: wrapped = with_revisions(revisions, wrapped) seen[oid] = wrapped return wrapped stream = _wrap(stream) for k, v in list(io.reads.items()): io.reads[k] = _wrap(v) for k, v in list(io.writes.items()): io.writes[k] = _wrap(v) return stream, io class MountEntry: """A mounted resource with command and op dispatch. Each mount has its own lookup tables for commands and ops. Different mounts of the same resource type can have different registered commands/ops. Resolution hierarchy (same for commands and ops): 1. (name, extension) -- filetype-specific 2. (name, None) -- resource-specific 3. general[name] -- general fallback """ def __init__( self, prefix: str, resource: BaseResource, mode: MountMode = MountMode.READ, consistency: ConsistencyPolicy = ConsistencyPolicy.LAZY, ) -> None: if not prefix.startswith("/"): raise ValueError(f"prefix must start with /: {prefix!r}") if not prefix.endswith("/"): raise ValueError(f"prefix must end with /: {prefix!r}") if "//" in prefix: raise ValueError(f"prefix must not contain //: {prefix!r}") self.prefix = prefix self.resource = resource self.mode = mode self.consistency = consistency self.cache_manager: CacheManager | None = None # Per-path revision pins installed at Workspace.load time. Read # functions consult these via the ``revision_for`` contextvar # lookup; on a hit, the backend GET pins to the recorded # revision so replay serves the exact bytes the agent saw. # Empty during normal runs; populated only by the snapshot # loader. self.revisions: dict[str, str] = {} self._cmds: dict[tuple, RegisteredCommand] = {} self._general_cmds: dict[str, RegisteredCommand] = {} self._cmd_specs: dict[str, CommandSpec] = {} self.command_safeguards: dict[str, CommandSafeguard] = {} self._ops: dict[tuple, RegisteredOp] = {} self._general_ops: dict[str, RegisteredOp] = {} # key: (cmd_name, target_resource_type) self._cross_cmds: dict[tuple, RegisteredCommand] = {} # ── command registration ────────────────────────── def register(self, cmd: RegisteredCommand) -> None: """Register a resource-specific command.""" key = (cmd.name, cmd.filetype) self._cmds[key] = cmd if cmd.spec is not None: self._cmd_specs[cmd.name] = cmd.spec def register_general( self, cmd: RegisteredCommand, ) -> None: """Register a general command (resource=None). General commands work on any resource (e.g. echo, pwd). They are the last fallback in resolve_command(). """ self._general_cmds[cmd.name] = cmd if cmd.spec is not None: self._cmd_specs[cmd.name] = cmd.spec def resolve_command( self, cmd_name: str, extension: str | None = None, ) -> RegisteredCommand | None: """Resolve command with fallback hierarchy. Lookup order: 1. (cmd_name, extension) -- filetype-specific 2. (cmd_name, None) -- resource-specific 3. general_cmds[cmd_name] -- general fallback """ if extension: cmd = self._cmds.get((cmd_name, extension)) if cmd is not None: return cmd cmd = self._cmds.get((cmd_name, None)) if cmd is not None: return cmd return self._general_cmds.get(cmd_name) def spec_for( self, cmd_name: str, ) -> CommandSpec | None: """Get the spec for a command name.""" return self._cmd_specs.get(cmd_name) def is_general_command(self, cmd_name: str) -> bool: """Whether `cmd_name` is registered as a general command here.""" return cmd_name in self._general_cmds def all_commands(self) -> list[RegisteredCommand]: """All registered commands (per-mount + general), deduped by name.""" seen: set[str] = set() out: list[RegisteredCommand] = [] for rc in self._cmds.values(): if rc.name in seen: continue seen.add(rc.name) out.append(rc) for rc in self._general_cmds.values(): if rc.name in seen: continue seen.add(rc.name) out.append(rc) return out def filetype_handlers( self, cmd_name: str, ) -> dict[str, Callable]: """Get filetype-specific command handlers. Example:: mount.register(generic_cat) # ("cat", None) mount.register(parquet_cat) # ("cat", ".parquet") mount.filetype_handlers("cat") # -> {".parquet": parquet_cat_fn} Args: cmd_name (str): command name, e.g. "cat". """ fns: dict[str, Callable] = {} for (name, ft), rc in self._cmds.items(): if name == cmd_name and ft is not None: if ft not in fns: fns[ft] = rc.fn return fns def register_fns(self, fns: list) -> None: """Register commands and ops from decorated functions. Args: fns (list): Functions decorated with @command and/or @op. Raises: ValueError: If a command/op's resource doesn't match this mount's resource. """ pname = self.resource.name for fn in fns: if hasattr(fn, "_registered_commands"): rcs = fn._registered_commands matching = [ rc for rc in rcs if rc.resource is None or rc.resource == pname ] if rcs and not matching: resources = sorted({rc.resource for rc in rcs}) raise ValueError( f"command {rcs[0].name!r} is for resource(s) " f"{resources!r}, not {pname!r}") for rc in matching: self.register(rc) if hasattr(fn, "_registered_ops"): ros = fn._registered_ops matching_ops = [ ro for ro in ros if ro.resource is None or ro.resource == pname ] if ros and not matching_ops: resources = sorted({ro.resource for ro in ros}) raise ValueError(f"op {ros[0].name!r} is for resource(s) " f"{resources!r}, not {pname!r}") for ro in matching_ops: self.register_op(ro) def unregister(self, names: list[str]) -> None: """Remove all commands and ops with the given names. Args: names (list[str]): Command/op names to remove. """ for name in names: keys = [k for k in self._cmds if k[0] == name] for k in keys: del self._cmds[k] self._general_cmds.pop(name, None) self._cmd_specs.pop(name, None) op_keys = [k for k in self._ops if k[0] == name] for k in op_keys: del self._ops[k] self._general_ops.pop(name, None) def commands(self) -> dict[str, list[str | None]]: """List registered commands grouped by filetype variants. Returns: dict[str, list[str | None]]: Command name to filetype list. """ result: dict[str, list[str | None]] = {} for (name, filetype) in self._cmds: result.setdefault(name, []).append(filetype) for name in self._general_cmds: result.setdefault(name, []) for name in result: result[name] = sorted(result[name], key=lambda x: (x is not None, x or "")) return dict(sorted(result.items())) def registered_ops(self) -> dict[str, list[str | None]]: """List registered ops grouped by filetype variants. Returns: dict[str, list[str | None]]: Op name to filetype list. """ result: dict[str, list[str | None]] = {} for (name, filetype) in self._ops: result.setdefault(name, []).append(filetype) for name in self._general_ops: result.setdefault(name, []) for name in result: result[name] = sorted(result[name], key=lambda x: (x is not None, x or "")) return dict(sorted(result.items())) # ── cross-mount registration ───────────────────── def register_cross( self, cmd: RegisteredCommand, target_resource_type: str, ) -> None: """Register a cross-mount command for a target. Example:: mount.register_cross(cp_cmd, "ram") # This mount can now cp to ram mounts Args: cmd: the cross-mount command. target_resource_type: e.g. "ram", "s3". """ key = (cmd.name, target_resource_type) self._cross_cmds[key] = cmd def resolve_cross( self, cmd_name: str, target_resource_type: str, ) -> RegisteredCommand | None: """Find a cross-mount command for a target.""" return self._cross_cmds.get((cmd_name, target_resource_type)) # ── op registration ─────────────────────────────── def register_op(self, op: RegisteredOp) -> None: """Register a resource-specific VFS op.""" key = (op.name, op.filetype) self._ops[key] = op def _resolve_cascade( self, name: str, extension: str | None, table: dict, general: dict, ) -> list: """Resolve with cascade: try filetype, resource, general. Returns list of matching entries to try in order. First non-None result wins. """ levels = [] if extension: entry = table.get((name, extension)) if entry is not None: levels.append(entry) entry = table.get((name, None)) if entry is not None: levels.append(entry) entry = general.get(name) if entry is not None: levels.append(entry) return levels # ── execution ───────────────────────────────────── async def execute_cmd( self, cmd_name: str, paths: list[PathSpec], texts: list[str], flag_kwargs: dict, *, stdin: ByteSource | None = None, cwd: str = "/", dispatch: Callable | None = None, session_id: str | None = None, env: dict[str, str] | None = None, exec_allowed: bool = True, ) -> tuple[ByteSource | None, IOResult]: """Execute a command on this mount's resource. Pure dispatch — flag parsing is done upstream in executor/command.py. This method just resolves the command handler and calls it. Args: cmd_name (str): command name. paths (list[PathSpec]): positional path args. texts (list[str]): positional text args. flag_kwargs (dict): parsed flags from upstream. stdin (ByteSource | None): stdin data. cwd (str): virtual cwd from session. """ extension = get_extension(paths[0].virtual) if paths else None handlers = self._resolve_cascade(cmd_name, extension, self._cmds, self._general_cmds) if not handlers: return None, IOResult( exit_code=127, stderr=(f"{cmd_name}: command not found".encode())) mount_prefix = self.prefix.rstrip("/") filetype_fns = self.filetype_handlers(cmd_name) is_filetype_cmd = extension is not None and (cmd_name, extension) in self._cmds paths = [ dataclasses.replace( p, resource_path=mount_key(p.virtual, mount_prefix)) if isinstance(p, PathSpec) else p for p in paths ] # Stamp this mount's backend key onto path-shaped flag values so # backend reads can address them: a single PathSpec (e.g. awk -f, # single grep -f) or a list of PathSpec (repeatable grep -f). # Everything else (bools, strings, list[str] like repeated -e) is # not a path and passes through unchanged. kw = {} for k, v in flag_kwargs.items(): if isinstance(v, PathSpec): kw[k] = dataclasses.replace(v, resource_path=mount_key( v.virtual, mount_prefix)) elif isinstance(v, list) and v and all( isinstance(item, PathSpec) for item in v): kw[k] = [ dataclasses.replace(item, resource_path=mount_key( item.virtual, mount_prefix)) for item in v ] else: kw[k] = v kw["index"] = self.resource.index kw["cwd"] = PathSpec( virtual=cwd, directory=cwd, resolved=False, resource_path=mount_key(cwd, mount_prefix), ) kw["filetype_fns"] = (filetype_fns if not is_filetype_cmd else None) if stdin is not None: kw["stdin"] = stdin if dispatch is not None: kw["dispatch"] = dispatch if session_id is not None: kw["session_id"] = session_id if env is not None: kw["env"] = env kw["exec_allowed"] = exec_allowed prev_prefix = push_mount_prefix(mount_prefix) revs_token = push_revisions(self.revisions or None) prev_manager = push_cache_manager(self.cache_manager) try: for cmd in handlers: if cmd.write and self.mode == MountMode.READ: return None, IOResult( exit_code=1, stderr=(f"{cmd_name}: read-only mount " f"at {self.prefix}".encode())) result = await cmd.fn(self.resource.accessor, paths, *texts, **kw) if result is not None: stream, io = _wrap_cmd_streams(result, mount_prefix, self.revisions or None) # TODO: hand back a finalization context separately # instead of stamping policy onto io.safeguard. io.safeguard = resolve_safeguard( cmd_name, cmd.safeguard, self.command_safeguards.get(cmd_name)) return stream, io return None, IOResult() finally: reset_revisions(revs_token) push_mount_prefix(prev_prefix) push_cache_manager(prev_manager) async def execute_op( self, op_name: str, path: str, *args, **kwargs, ) -> Any: """Execute a VFS op on this mount's resource. Tries filetype-specific first, then resource-specific. First non-None result wins. Args: op_name (str): operation name (e.g. "read", "stat"). path (str): virtual path. """ filetype = get_extension(path) levels = self._resolve_cascade(op_name, filetype, self._ops, self._general_ops) if not levels: raise AttributeError(f"{self.resource.name}: " f"no op {op_name!r}") if self.mode == MountMode.READ and any(o.write for o in levels): raise PermissionError(f"mount {self.prefix!r} is read-only") mount_prefix = self.prefix.rstrip("/") scope = PathSpec( virtual=path, directory=path.rsplit("/", 1)[0] or "/", resource_path=mount_key(path, mount_prefix), ) kwargs.setdefault("index", self.resource.index) op_override = self.command_safeguards.get(op_name) op_timeout = (op_override.timeout_seconds if op_override is not None else None) prev_prefix = push_mount_prefix(mount_prefix) revs_token = push_revisions(self.revisions or None) try: for op in levels: result = op.fn(self.resource.accessor, scope, *args, **kwargs) if inspect.isawaitable(result): result = await run_with_timeout(result, op_timeout, op_name) if result is not None: return await apply_op_safeguard(result, op_override) return None finally: reset_revisions(revs_token) push_mount_prefix(prev_prefix)